Type guard for t.discriminatedUnion

[PatrickRose]

I'm attempting to use discriminated unions for some of the new feature we're doing. I'd written the following code but I get a TS error

import {z} from 'zod'

export function ajaxReturnSchema<T extends z.ZodTypeAny>(schema: T): z.ZodObject<{ code: z.ZodNumber, msg: z.ZodNullable<z.ZodString>, data: T  }> {
    return z.object({
        code: z.number(),
        msg: z.string().nullable(),
        data: schema
    });
}

export function handleDiscriminatedAjaxReturn<Success extends z.ZodTypeAny, SuccessReturn, Fail extends z.ZodTypeAny, FailReturn>(
    data: unknown,
    success: Success, 
    fail: Fail, 
    successCallback: (success: z.infer<Success>) => SuccessReturn, 
    failureCallback: (fail: z.infer<Fail>) => FailReturn
): SuccessReturn | FailReturn {
    const schema = ajaxReturnSchema(
        z.discriminatedUnion(
            'status',
            [
                z.object({status: z.literal('success'), result: success}),
                z.object({status: z.literal('failure'), error: fail})
            ]
        )
    );

    const result = schema.safeParse(data);

    if (!result.success) {
        throw new Error();
    }

    const parsedData = result.data.data;

    if (parsedData.status == 'success') {
        return successCallback(parsedData.result);
    } else {
        return failureCallback(parsedData.error)
    }
}

const fail = z.object({testFail: z.string()});
type fail = z.infer<typeof fail>

const failCallback = (fail: fail) => {
    return fail.testFail;
}

const success = z.object({test: z.number()});
type Success = z.infer<typeof success>

const successCallback = (success: Success) => {
    return success.test;
}

const result = handleDiscriminatedAjaxReturn('', success, fail, successCallback, failCallback);

However, I get told:

• parsedData is possibly undefined (which I wasn't expecting)
• .result / .error don't exist on the type, because the compiler still treats it as {'status: 'success', result: z.infer<Success>} | {'status: 'failure', result: z.infer<Failure>}, even though there's the if conditional