Skip to content

Commit 53006b7

Browse files
easonLiangWorldedtecheasonLiangWorldedtech
committed
fix(kimi-code): keep the OAuth abort forwarder attached until the body settles
An abort that landed after the response headers rejected the caller's wait but left the underlying POST open: postForm detached the caller-signal forwarder as soon as the headers arrived, so readBodyWithSignal cancelled only the wrapper and a replacement refresh could start on top of a response nobody was waiting for. postForm now hands back a settle hook and each caller detaches it after its own body read. Two tests pin the abort reaching the fetch signal, one on a stalled success body and one on a stalled OAuth error body. The zai spec's console.error restore is moved to span the awaited captureError, and a second test pins that a resolving mocked create leaves no spy behind.
1 parent 6762b55 commit 53006b7

3 files changed

Lines changed: 162 additions & 33 deletions

File tree

‎src/api/providers/__tests__/zai.spec.ts‎

Lines changed: 39 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -714,22 +714,53 @@ describe("ZAiHandler", () => {
714714
mockCreate.mockRejectedValueOnce(new Error("boom"))
715715
const errorLog = vitest.spyOn(console, "error").mockImplementation(() => {})
716716

717-
const result = await captureError(
718-
(async () => {
719-
for await (const _ of thinkingHandler.createMessage("system prompt", [])) {
720-
// consume
721-
}
722-
})(),
723-
)
724-
717+
// captureError throws when the operation resolves instead of rejecting - the very contract
718+
// violation this test exists to catch - so the restoration has to span the await: a restoring
719+
// finally that starts after it leaks the console.error spy into whichever test runs next.
720+
let result: Error
725721
try {
722+
result = await captureError(
723+
(async () => {
724+
for await (const _ of thinkingHandler.createMessage("system prompt", [])) {
725+
// consume
726+
}
727+
})(),
728+
)
726729
expect(result.message).toBe("Z.ai completion error: boom")
727730
expect(errorLog).toHaveBeenCalledTimes(1)
728731
} finally {
729732
errorLog.mockRestore()
730733
}
731734
})
732735

736+
it("createMessage leaves no console.error spy installed when the mocked create resolves", async () => {
737+
const thinkingHandler = new ZAiHandler({
738+
apiModelId: "glm-4.7",
739+
zaiApiKey: "test-zai-api-key",
740+
zaiApiLine: "international_coding",
741+
})
742+
// The resolving shape: an empty async iterable satisfies the stream loop and rejects nothing,
743+
// so the only failure left is captureError's own "expected a rejection" assertion.
744+
mockCreate.mockResolvedValueOnce((async function* () {})())
745+
const errorLog = vitest.spyOn(console, "error").mockImplementation(() => {})
746+
let failure: Error | undefined
747+
try {
748+
failure = await captureError(
749+
(async () => {
750+
for await (const _ of thinkingHandler.createMessage("system prompt", [])) {
751+
// consume
752+
}
753+
})(),
754+
)
755+
} catch (error) {
756+
failure = error instanceof Error ? error : new Error(String(error))
757+
} finally {
758+
errorLog.mockRestore()
759+
}
760+
expect(failure?.message).toBe("Expected the operation to reject")
761+
expect(vitest.isMockFunction(console.error), "the spy must not outlive its own test").toBe(false)
762+
})
763+
733764
it("createMessage should normalize an abort error raised during stream iteration on the thinking path", async () => {
734765
const thinkingHandler = new ZAiHandler({
735766
apiModelId: "glm-4.7",

‎src/integrations/kimi-code/__tests__/oauth.spec.ts‎

Lines changed: 68 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1028,6 +1028,74 @@ describe("KimiCodeOAuthManager", () => {
10281028
expect(bodyReads).toEqual(["json"])
10291029
})
10301030

1031+
it("cancels the underlying response when the caller aborts during the token body read", async () => {
1032+
// The wrapper rejection alone is not cancellation: the fetch signal that carries the POST must
1033+
// still be armed while the body is pending, or a replacement refresh can start on top of a
1034+
// response this caller has already given up on.
1035+
let fetchedSignal: AbortSignal | undefined
1036+
const bodyReads: string[] = []
1037+
const bodyPending = new Promise<unknown>(() => {})
1038+
vi.spyOn(globalThis, "fetch").mockImplementation(((_url: unknown, init?: RequestInit) => {
1039+
fetchedSignal = init?.signal as AbortSignal | undefined
1040+
return Promise.resolve({
1041+
ok: true,
1042+
status: 200,
1043+
json: () => {
1044+
bodyReads.push("json")
1045+
return bodyPending
1046+
},
1047+
} as unknown as Response)
1048+
}) as typeof fetch)
1049+
const controller = new AbortController()
1050+
const credentials: KimiCodeCredentials = {
1051+
type: "kimi-code",
1052+
accessToken: "old",
1053+
refreshToken: "refresh",
1054+
expiresAt: 0,
1055+
}
1056+
const pending = refreshKimiCodeAccessToken(credentials, controller.signal)
1057+
await vi.waitFor(() => expect(bodyReads).toEqual(["json"]))
1058+
expect(fetchedSignal?.aborted, "the POST signal stays live while the body is pending").toBe(false)
1059+
controller.abort()
1060+
await expect(pending).rejects.toMatchObject({ name: "AbortError" })
1061+
expect(fetchedSignal?.aborted, "the abort reached the underlying response").toBe(true)
1062+
})
1063+
1064+
it("cancels the underlying response when the caller aborts during a stalled OAuth error body", async () => {
1065+
// The failure path reads a body too. An error body that stalls must be cancelled the same way,
1066+
// or the rejected refresh still holds the response open after the caller is gone.
1067+
let fetchedSignal: AbortSignal | undefined
1068+
const bodyReads: string[] = []
1069+
const bodyPending = new Promise<unknown>(() => {})
1070+
vi.spyOn(globalThis, "fetch").mockImplementation(((_url: unknown, init?: RequestInit) => {
1071+
fetchedSignal = init?.signal as AbortSignal | undefined
1072+
return Promise.resolve({
1073+
ok: false,
1074+
status: 400,
1075+
statusText: "Bad Request",
1076+
text: () => {
1077+
bodyReads.push("text")
1078+
return bodyPending
1079+
},
1080+
} as unknown as Response)
1081+
}) as typeof fetch)
1082+
const controller = new AbortController()
1083+
const credentials: KimiCodeCredentials = {
1084+
type: "kimi-code",
1085+
accessToken: "old",
1086+
refreshToken: "refresh",
1087+
expiresAt: 0,
1088+
}
1089+
const pending = refreshKimiCodeAccessToken(credentials, controller.signal)
1090+
await vi.waitFor(() => expect(bodyReads).toEqual(["text"]))
1091+
expect(fetchedSignal?.aborted, "the POST signal stays live while the error body is pending").toBe(false)
1092+
controller.abort()
1093+
await expect(pending).rejects.toMatchObject({
1094+
name: "AbortError",
1095+
message: "The Kimi Code OAuth request was aborted",
1096+
})
1097+
expect(fetchedSignal?.aborted, "the abort reached the underlying error-body response").toBe(true)
1098+
})
10311099
it("never aborts a replacement refresh on behalf of a waiter enrolled in the previous one", async () => {
10321100
const { context, values } = createContext()
10331101
values.set(

‎src/integrations/kimi-code/oauth.ts‎

Lines changed: 55 additions & 25 deletions
Original file line numberDiff line numberDiff line change
@@ -79,7 +79,19 @@ function createOAuthAbortError(): Error {
7979
return error
8080
}
8181

82-
async function postForm(endpoint: string, values: Record<string, string>, signal?: AbortSignal): Promise<Response> {
82+
/** A form POST whose cancellation contract outlives the response headers: `settle` detaches the
83+
* caller's abort forwarder once the body is consumed, so an abort during a body read cancels the
84+
* underlying response too. An abort detaches the forwarder by itself ({ once: true }). */
85+
interface CancellingPost {
86+
response: Response
87+
settle: () => void
88+
}
89+
90+
async function postForm(
91+
endpoint: string,
92+
values: Record<string, string>,
93+
signal?: AbortSignal,
94+
): Promise<CancellingPost> {
8395
const controller = new AbortController()
8496
const forwardAbort = () => controller.abort(signal?.reason)
8597
if (signal?.aborted) forwardAbort()
@@ -88,8 +100,11 @@ async function postForm(endpoint: string, values: Record<string, string>, signal
88100
() => controller.abort(new Error("Kimi Code OAuth request timed out")),
89101
OAUTH_REQUEST_TIMEOUT_MS,
90102
)
103+
const settle = () => {
104+
signal?.removeEventListener("abort", forwardAbort)
105+
}
91106
try {
92-
return await fetch(endpoint, {
107+
const response = await fetch(endpoint, {
93108
method: "POST",
94109
headers: {
95110
Accept: "application/json",
@@ -98,7 +113,11 @@ async function postForm(endpoint: string, values: Record<string, string>, signal
98113
body: new URLSearchParams(values).toString(),
99114
signal: controller.signal,
100115
})
116+
return { response, settle }
101117
} catch (error) {
118+
// Nothing will read a body from a POST that never produced a response, so the forwarder
119+
// is detached on this path and nothing else.
120+
settle()
102121
// Classify from the CALLER signal: only a rejection that came from the caller's own abort
103122
// is normalized. The caller's reason can be an ordinary Error and fetch forwards it
104123
// verbatim, and an independent OAuth/network failure that merely races the abort must be
@@ -113,7 +132,6 @@ async function postForm(endpoint: string, values: Record<string, string>, signal
113132
throw error
114133
} finally {
115134
clearTimeout(timeout)
116-
signal?.removeEventListener("abort", forwardAbort)
117135
}
118136
}
119137
/** Consume a response body under the caller's signal: an abort while the body is still
@@ -148,17 +166,21 @@ async function readOAuthError(response: Response, signal?: AbortSignal): Promise
148166
}
149167

150168
export async function requestDeviceAuthorization(signal?: AbortSignal) {
151-
const response = await postForm(
169+
const { response, settle } = await postForm(
152170
KIMI_CODE_OAUTH_CONFIG.deviceAuthorizationEndpoint,
153171
{ client_id: KIMI_CODE_OAUTH_CONFIG.clientId },
154172
signal,
155173
)
156-
if (!response.ok) throw await readOAuthError(response, signal)
157-
return deviceAuthorizationSchema.parse(await readBodyWithSignal(() => response.json(), signal))
174+
try {
175+
if (!response.ok) throw await readOAuthError(response, signal)
176+
return deviceAuthorizationSchema.parse(await readBodyWithSignal(() => response.json(), signal))
177+
} finally {
178+
settle()
179+
}
158180
}
159181

160182
async function requestDeviceToken(deviceCode: string, signal?: AbortSignal): Promise<KimiCodeCredentials> {
161-
const response = await postForm(
183+
const { response, settle } = await postForm(
162184
KIMI_CODE_OAUTH_CONFIG.tokenEndpoint,
163185
{
164186
grant_type: KIMI_CODE_OAUTH_CONFIG.deviceGrantType,
@@ -167,15 +189,19 @@ async function requestDeviceToken(deviceCode: string, signal?: AbortSignal): Pro
167189
},
168190
signal,
169191
)
170-
if (!response.ok) throw await readOAuthError(response, signal)
171-
const tokens = tokenResponseSchema.parse(await readBodyWithSignal(() => response.json(), signal))
172-
if (!tokens.refresh_token) throw new Error("Kimi Code OAuth did not return a refresh token")
173-
return {
174-
type: providerIdentifiers.kimiCode,
175-
accessToken: tokens.access_token,
176-
refreshToken: tokens.refresh_token,
177-
expiresAt: Date.now() + tokens.expires_in * 1000,
178-
tokenType: tokens.token_type,
192+
try {
193+
if (!response.ok) throw await readOAuthError(response, signal)
194+
const tokens = tokenResponseSchema.parse(await readBodyWithSignal(() => response.json(), signal))
195+
if (!tokens.refresh_token) throw new Error("Kimi Code OAuth did not return a refresh token")
196+
return {
197+
type: providerIdentifiers.kimiCode,
198+
accessToken: tokens.access_token,
199+
refreshToken: tokens.refresh_token,
200+
expiresAt: Date.now() + tokens.expires_in * 1000,
201+
tokenType: tokens.token_type,
202+
}
203+
} finally {
204+
settle()
179205
}
180206
}
181207

@@ -185,7 +211,7 @@ export async function refreshKimiCodeAccessToken(
185211
): Promise<KimiCodeCredentials> {
186212
// The caller's signal aborts the token POST itself, not just the caller's wait: a cancelled
187213
// task must not keep a refresh request (and its rotated refresh token) in flight.
188-
const response = await postForm(
214+
const { response, settle } = await postForm(
189215
KIMI_CODE_OAUTH_CONFIG.tokenEndpoint,
190216
{
191217
grant_type: "refresh_token",
@@ -194,14 +220,18 @@ export async function refreshKimiCodeAccessToken(
194220
},
195221
signal,
196222
)
197-
if (!response.ok) throw await readOAuthError(response, signal)
198-
const tokens = tokenResponseSchema.parse(await readBodyWithSignal(() => response.json(), signal))
199-
return {
200-
type: providerIdentifiers.kimiCode,
201-
accessToken: tokens.access_token,
202-
refreshToken: tokens.refresh_token ?? credentials.refreshToken,
203-
expiresAt: Date.now() + tokens.expires_in * 1000,
204-
tokenType: tokens.token_type ?? credentials.tokenType,
223+
try {
224+
if (!response.ok) throw await readOAuthError(response, signal)
225+
const tokens = tokenResponseSchema.parse(await readBodyWithSignal(() => response.json(), signal))
226+
return {
227+
type: providerIdentifiers.kimiCode,
228+
accessToken: tokens.access_token,
229+
refreshToken: tokens.refresh_token ?? credentials.refreshToken,
230+
expiresAt: Date.now() + tokens.expires_in * 1000,
231+
tokenType: tokens.token_type ?? credentials.tokenType,
232+
}
233+
} finally {
234+
settle()
205235
}
206236
}
207237

0 commit comments

Comments
 (0)