mirror of
https://github.com/Gitlawb/openclaude.git
synced 2026-08-24 02:34:15 -05:00
fix(typecheck): recreate missing CLI Transport interface (#1581)
* fix(typecheck): recreate missing CLI Transport interface * fix(transports): implement async close in CLI transports * fix(transports): harden async close cleanup * fix: address async transport close review feedback * test: isolate environment-sensitive suites * fix(transports): drain uploader after close failure * fix(transports): type guard CCR stream events * test: remove unused auto-compact fixture helper
This commit is contained in:
Bogdan
@@ -490,7 +490,7 @@ export async function initEnvLessBridgeCore(
|
||||
flushGate.start()
|
||||
try {
|
||||
const seq = transport.getLastSequenceNum()
|
||||
transport.close()
|
||||
await transport.close()
|
||||
transport = await createV2ReplTransport({
|
||||
sessionUrl: buildCCRv2SdkUrl(fresh.api_base_url, sessionId),
|
||||
ingressToken: fresh.worker_jwt,
|
||||
@@ -506,7 +506,7 @@ export async function initEnvLessBridgeCore(
|
||||
// Teardown fired during the async createV2ReplTransport window.
|
||||
// Don't wire/connect/schedule — we'd re-arm timers after cancelAll()
|
||||
// and fire onInboundMessage into a torn-down bridge.
|
||||
transport.close()
|
||||
await transport.close()
|
||||
return
|
||||
}
|
||||
wireTransportCallbacks()
|
||||
@@ -717,7 +717,14 @@ export async function initEnvLessBridgeCore(
|
||||
}
|
||||
}
|
||||
|
||||
transport.close()
|
||||
try {
|
||||
await transport.close()
|
||||
} catch (err) {
|
||||
logForDebugging(
|
||||
`[remote-bridge] Transport close threw during teardown: ${errorMessage(err)}`,
|
||||
{ level: 'error' },
|
||||
)
|
||||
}
|
||||
|
||||
const archiveStatus: ArchiveTelemetryStatus =
|
||||
status === 'no_token'
|
||||
|
||||
+39
-14
@@ -626,6 +626,20 @@ export async function initBridgeCore(
|
||||
}
|
||||
}
|
||||
|
||||
async function closeTransportBestEffort(
|
||||
transportToClose: ReplBridgeTransport,
|
||||
reason: string,
|
||||
): Promise<void> {
|
||||
try {
|
||||
await transportToClose.close()
|
||||
} catch (err) {
|
||||
logForDebugging(
|
||||
`[bridge:repl] Transport close threw ${reason}: ${errorMessage(err)}`,
|
||||
{ level: 'error' },
|
||||
)
|
||||
}
|
||||
}
|
||||
|
||||
async function doReconnect(): Promise<boolean> {
|
||||
environmentRecreations++
|
||||
// Invalidate any in-flight v2 handshake — the environment is being
|
||||
@@ -652,7 +666,7 @@ export async function initBridgeCore(
|
||||
if (seq > lastTransportSequenceNum) {
|
||||
lastTransportSequenceNum = seq
|
||||
}
|
||||
transport.close()
|
||||
await closeTransportBestEffort(transport, 'during reconnect')
|
||||
transport = null
|
||||
}
|
||||
// Transport is gone — wake the poll loop out of its at-capacity
|
||||
@@ -1055,7 +1069,7 @@ export async function initBridgeCore(
|
||||
// forwarding prompts → ~25-min dead window observed in daemon logs.
|
||||
// Kill the transport + work state so isAtCapacity()=false; the loop
|
||||
// fast-polls and picks up the server's re-dispatched work in seconds.
|
||||
onHeartbeatFatal: (err: BridgeFatalError) => {
|
||||
onHeartbeatFatal: async (err: BridgeFatalError) => {
|
||||
logForDebugging(
|
||||
`[bridge:repl] heartbeatWork fatal (status=${err.status}) — tearing down work item for fast re-dispatch`,
|
||||
)
|
||||
@@ -1064,7 +1078,7 @@ export async function initBridgeCore(
|
||||
if (seq > lastTransportSequenceNum) {
|
||||
lastTransportSequenceNum = seq
|
||||
}
|
||||
transport.close()
|
||||
await closeTransportBestEffort(transport, 'after heartbeat fatal')
|
||||
transport = null
|
||||
}
|
||||
flushGate.drop()
|
||||
@@ -1094,7 +1108,7 @@ export async function initBridgeCore(
|
||||
}
|
||||
return { environmentId, environmentSecret }
|
||||
},
|
||||
onWorkReceived: (
|
||||
onWorkReceived: async (
|
||||
workSessionId: string,
|
||||
ingressToken: string,
|
||||
workId: string,
|
||||
@@ -1199,7 +1213,10 @@ export async function initBridgeCore(
|
||||
if (oldSeq > lastTransportSequenceNum) {
|
||||
lastTransportSequenceNum = oldSeq
|
||||
}
|
||||
oldTransport.close()
|
||||
await closeTransportBestEffort(
|
||||
oldTransport,
|
||||
'while replacing work transport',
|
||||
)
|
||||
}
|
||||
// Reset flush state — the old flush (if any) is no longer relevant.
|
||||
// Preserve pending messages so they're drained after the new
|
||||
@@ -1410,13 +1427,16 @@ export async function initBridgeCore(
|
||||
sessionId: workSessionId,
|
||||
initialSequenceNum: lastTransportSequenceNum,
|
||||
}).then(
|
||||
t => {
|
||||
async t => {
|
||||
// Teardown started while registerWorker was in flight. Teardown
|
||||
// saw transport === null and skipped close(); installing now
|
||||
// would leak CCRClient heartbeat timers and reset
|
||||
// teardownStarted via wireTransport's side effects.
|
||||
if (pollController.signal.aborted) {
|
||||
t.close()
|
||||
await closeTransportBestEffort(
|
||||
t,
|
||||
'while discarding aborted CCR v2 transport',
|
||||
)
|
||||
return
|
||||
}
|
||||
// onWorkReceived may have fired again while registerWorker()
|
||||
@@ -1429,7 +1449,10 @@ export async function initBridgeCore(
|
||||
logForDebugging(
|
||||
`[bridge:repl] CCR v2: discarding stale handshake gen=${thisGen} current=${v2Generation}`,
|
||||
)
|
||||
t.close()
|
||||
await closeTransportBestEffort(
|
||||
t,
|
||||
'while discarding stale CCR v2 transport',
|
||||
)
|
||||
return
|
||||
}
|
||||
wireTransport(t)
|
||||
@@ -1668,8 +1691,10 @@ export async function initBridgeCore(
|
||||
// log their own success/failure internally.
|
||||
await Promise.all([stopWorkP, archiveSession(currentSessionId)])
|
||||
|
||||
teardownTransport?.close()
|
||||
logForDebugging('[bridge:repl] Teardown: transport closed')
|
||||
if (teardownTransport) {
|
||||
await closeTransportBestEffort(teardownTransport, 'during teardown')
|
||||
logForDebugging('[bridge:repl] Teardown: transport closed')
|
||||
}
|
||||
|
||||
await api.deregisterEnvironment(environmentId).catch((err: unknown) => {
|
||||
logForDebugging(
|
||||
@@ -1892,7 +1917,7 @@ async function startWorkPollLoop({
|
||||
ingressToken: string,
|
||||
workId: string,
|
||||
useCodeSessions: boolean,
|
||||
) => void
|
||||
) => Promise<void>
|
||||
/** Called when the environment has been deleted. Returns new credentials or null. */
|
||||
onEnvironmentLost?: () => Promise<{
|
||||
environmentId: string
|
||||
@@ -1935,7 +1960,7 @@ async function startWorkPollLoop({
|
||||
* ~10-minute dead window before recovery). When omitted, falls back to
|
||||
* the backoff sleep to avoid a tight poll+heartbeat loop.
|
||||
*/
|
||||
onHeartbeatFatal?: (err: BridgeFatalError) => void
|
||||
onHeartbeatFatal?: (err: BridgeFatalError) => Promise<void>
|
||||
}): Promise<void> {
|
||||
const MAX_ENVIRONMENT_RECREATIONS = 3
|
||||
|
||||
@@ -2063,7 +2088,7 @@ async function startWorkPollLoop({
|
||||
// for the server's re-dispatched work item. Without
|
||||
// the hook, backoff to avoid tight poll+heartbeat loop.
|
||||
if (onHeartbeatFatal) {
|
||||
onHeartbeatFatal(err)
|
||||
await onHeartbeatFatal(err)
|
||||
logForDebugging(
|
||||
`[bridge:repl:heartbeat] Fatal (status=${err.status}), work state cleared — fast-polling for re-dispatch`,
|
||||
)
|
||||
@@ -2197,7 +2222,7 @@ async function startWorkPollLoop({
|
||||
continue
|
||||
}
|
||||
|
||||
onWorkReceived(
|
||||
await onWorkReceived(
|
||||
workSessionId,
|
||||
secret.session_ingress_token,
|
||||
work.id,
|
||||
|
||||
@@ -23,7 +23,7 @@ import { registerWorker } from './workSecret.js'
|
||||
export type ReplBridgeTransport = {
|
||||
write(message: StdoutMessage): Promise<void>
|
||||
writeBatch(messages: StdoutMessage[]): Promise<void>
|
||||
close(): void
|
||||
close(): Promise<void>
|
||||
isConnectedStatus(): boolean
|
||||
getStateLabel(): string
|
||||
setOnData(callback: (data: string) => void): void
|
||||
@@ -210,20 +210,8 @@ export async function createV2ReplTransport(opts: {
|
||||
logForDebugging(
|
||||
'[bridge:repl] CCR v2: epoch superseded (409) — closing for poll-loop recovery',
|
||||
)
|
||||
// Close resources in a try block so the throw always executes.
|
||||
// If ccr.close() or sse.close() throw, we still need to unwind
|
||||
// the caller (request()) — otherwise handleEpochMismatch's `never`
|
||||
// return type is violated at runtime and control falls through.
|
||||
try {
|
||||
ccr.close()
|
||||
sse.close()
|
||||
onCloseCb?.(4090)
|
||||
} catch (closeErr: unknown) {
|
||||
logForDebugging(
|
||||
`[bridge:repl] CCR v2: error during epoch-mismatch cleanup: ${errorMessage(closeErr)}`,
|
||||
{ level: 'error' },
|
||||
)
|
||||
}
|
||||
closeResourcesBestEffort('epoch-mismatch cleanup')
|
||||
onCloseCb?.(4090)
|
||||
// Don't return — the calling request() code continues after the 409
|
||||
// branch, so callers see the logged warning and a false return. We
|
||||
// throw to unwind; the uploaders catch it as a send failure.
|
||||
@@ -267,6 +255,30 @@ export async function createV2ReplTransport(opts: {
|
||||
let ccrInitialized = false
|
||||
let closed = false
|
||||
|
||||
function closeResources(): Promise<void> {
|
||||
closed = true
|
||||
return Promise.all([ccr.close(), sse.close()]).then(() => {})
|
||||
}
|
||||
|
||||
function closeResourcesBestEffort(reason: string): void {
|
||||
void closeResources().catch((closeErr: unknown) => {
|
||||
logForDebugging(
|
||||
`[bridge:repl] CCR v2: error during ${reason}: ${errorMessage(closeErr)}`,
|
||||
{ level: 'error' },
|
||||
)
|
||||
})
|
||||
}
|
||||
|
||||
function closeCcrBestEffort(reason: string): void {
|
||||
closed = true
|
||||
void ccr.close().catch((closeErr: unknown) => {
|
||||
logForDebugging(
|
||||
`[bridge:repl] CCR v2: error during ${reason}: ${errorMessage(closeErr)}`,
|
||||
{ level: 'error' },
|
||||
)
|
||||
})
|
||||
}
|
||||
|
||||
return {
|
||||
write(msg) {
|
||||
return ccr.writeEvent(msg)
|
||||
@@ -281,11 +293,7 @@ export async function createV2ReplTransport(opts: {
|
||||
await ccr.writeEvent(m)
|
||||
}
|
||||
},
|
||||
close() {
|
||||
closed = true
|
||||
ccr.close()
|
||||
sse.close()
|
||||
},
|
||||
close: closeResources,
|
||||
isConnectedStatus() {
|
||||
// Write-readiness, not read-readiness — replBridge checks this
|
||||
// before calling writeBatch. SSE open state is orthogonal.
|
||||
@@ -309,7 +317,7 @@ export async function createV2ReplTransport(opts: {
|
||||
// heartbeat timer before notifying replBridge. (sse.close() doesn't
|
||||
// invoke this, so the epoch-mismatch path above isn't double-firing.)
|
||||
sse.setOnClose(code => {
|
||||
ccr.close()
|
||||
closeCcrBestEffort('SSE close cleanup')
|
||||
cb(code ?? 4092)
|
||||
})
|
||||
},
|
||||
@@ -360,8 +368,7 @@ export async function createV2ReplTransport(opts: {
|
||||
// so the poll loop can retry on the next work dispatch.
|
||||
// Without this callback, replBridge never learns the transport
|
||||
// failed to initialize and sits with transport === null forever.
|
||||
ccr.close()
|
||||
sse.close()
|
||||
closeResourcesBestEffort('initialize failure cleanup')
|
||||
onCloseCb?.(4091) // 4091 = init failure, distinguishable from 4090 epoch mismatch
|
||||
},
|
||||
)
|
||||
|
||||
@@ -0,0 +1,99 @@
|
||||
import { expect, test } from 'bun:test'
|
||||
import { BridgeFatalError } from './bridgeApi.js'
|
||||
import { DEFAULT_POLL_CONFIG } from './pollConfigDefaults.js'
|
||||
import { _startWorkPollLoopForTesting } from './replBridge.js'
|
||||
import type { ReplBridgeTransport } from './replBridgeTransport.js'
|
||||
import type { BridgeApiClient } from './types.js'
|
||||
|
||||
type AssertPromiseVoid<T extends Promise<void>> = T
|
||||
|
||||
type _CloseReturnsPromise = AssertPromiseVoid<
|
||||
ReturnType<ReplBridgeTransport['close']>
|
||||
>
|
||||
|
||||
test('work poll loop waits for heartbeat-fatal cleanup before fast-polling again', async () => {
|
||||
const abort = new AbortController()
|
||||
let pollCount = 0
|
||||
let resolveCleanup: (() => void) | undefined
|
||||
const cleanupStarted = deferred<void>()
|
||||
const secondPollStarted = deferred<void>()
|
||||
const cleanupReleased = new Promise<void>(resolve => {
|
||||
resolveCleanup = resolve
|
||||
})
|
||||
|
||||
const api = {
|
||||
pollForWork: async () => {
|
||||
pollCount += 1
|
||||
if (pollCount === 2) {
|
||||
secondPollStarted.resolve()
|
||||
abort.abort()
|
||||
}
|
||||
return null
|
||||
},
|
||||
heartbeatWork: async () => {
|
||||
throw new BridgeFatalError('work item gone', 404)
|
||||
},
|
||||
} as unknown as BridgeApiClient
|
||||
|
||||
let atCapacity = true
|
||||
const loop = _startWorkPollLoopForTesting({
|
||||
api,
|
||||
getCredentials: () => ({
|
||||
environmentId: 'env-1',
|
||||
environmentSecret: 'secret-1',
|
||||
}),
|
||||
signal: abort.signal,
|
||||
isAtCapacity: () => atCapacity,
|
||||
capacitySignal: createCapacitySignal,
|
||||
getHeartbeatInfo: () => ({
|
||||
environmentId: 'env-1',
|
||||
workId: 'work-1',
|
||||
sessionToken: 'token-1',
|
||||
}),
|
||||
getPollIntervalConfig: () => ({
|
||||
...DEFAULT_POLL_CONFIG,
|
||||
poll_interval_ms_not_at_capacity: 1,
|
||||
poll_interval_ms_at_capacity: 10_000,
|
||||
non_exclusive_heartbeat_interval_ms: 1,
|
||||
reclaim_older_than_ms: 0,
|
||||
}),
|
||||
onWorkReceived: async () => {},
|
||||
onHeartbeatFatal: async () => {
|
||||
atCapacity = false
|
||||
cleanupStarted.resolve()
|
||||
await cleanupReleased
|
||||
},
|
||||
})
|
||||
|
||||
await cleanupStarted.promise
|
||||
await Promise.resolve()
|
||||
await Promise.resolve()
|
||||
expect(pollCount).toBe(1)
|
||||
|
||||
resolveCleanup?.()
|
||||
await secondPollStarted.promise
|
||||
await loop
|
||||
expect(pollCount).toBe(2)
|
||||
})
|
||||
|
||||
function createCapacitySignal(): {
|
||||
signal: AbortSignal
|
||||
cleanup: () => void
|
||||
} {
|
||||
const controller = new AbortController()
|
||||
return {
|
||||
signal: controller.signal,
|
||||
cleanup: () => {},
|
||||
}
|
||||
}
|
||||
|
||||
function deferred<T>(): {
|
||||
promise: Promise<T>
|
||||
resolve: (value: T | PromiseLike<T>) => void
|
||||
} {
|
||||
let resolve!: (value: T | PromiseLike<T>) => void
|
||||
const promise = new Promise<T>(res => {
|
||||
resolve = res
|
||||
})
|
||||
return { promise, resolve }
|
||||
}
|
||||
+2
-2
@@ -244,12 +244,12 @@ export class RemoteIO extends StructuredIO {
|
||||
/**
|
||||
* Clean up connections gracefully
|
||||
*/
|
||||
close(): void {
|
||||
async close(): Promise<void> {
|
||||
if (this.keepAliveTimer) {
|
||||
clearInterval(this.keepAliveTimer)
|
||||
this.keepAliveTimer = null
|
||||
}
|
||||
this.transport.close()
|
||||
await this.transport.close()
|
||||
this.inputStream.end()
|
||||
}
|
||||
}
|
||||
|
||||
@@ -0,0 +1,207 @@
|
||||
import { afterEach, beforeEach, describe, expect, jest, mock, test } from 'bun:test'
|
||||
import type { StdoutMessage } from 'src/entrypoints/sdk/controlTypes.js'
|
||||
|
||||
type AxiosPost = (
|
||||
url: string,
|
||||
data?: unknown,
|
||||
config?: unknown,
|
||||
) => Promise<{ status: number }>
|
||||
|
||||
let postImpl: AxiosPost = async () => ({ status: 200 })
|
||||
|
||||
mock.module('axios', () => ({
|
||||
default: {
|
||||
post: (...args: Parameters<AxiosPost>) => postImpl(...args),
|
||||
},
|
||||
}))
|
||||
|
||||
describe('HybridTransport close', () => {
|
||||
let originalSessionAccessToken: string | undefined
|
||||
|
||||
beforeEach(() => {
|
||||
originalSessionAccessToken = process.env.CLAUDE_CODE_SESSION_ACCESS_TOKEN
|
||||
process.env.CLAUDE_CODE_SESSION_ACCESS_TOKEN = 'test-token'
|
||||
})
|
||||
|
||||
afterEach(() => {
|
||||
if (originalSessionAccessToken === undefined) {
|
||||
delete process.env.CLAUDE_CODE_SESSION_ACCESS_TOKEN
|
||||
} else {
|
||||
process.env.CLAUDE_CODE_SESSION_ACCESS_TOKEN = originalSessionAccessToken
|
||||
}
|
||||
postImpl = async () => ({ status: 200 })
|
||||
jest.restoreAllMocks()
|
||||
})
|
||||
|
||||
test('drains buffered stream events before closing the uploader', async () => {
|
||||
const posts: Array<{ url: string; data: unknown }> = []
|
||||
postImpl = async (url, data) => {
|
||||
posts.push({ url, data })
|
||||
return { status: 200 }
|
||||
}
|
||||
const transport = await createTransport()
|
||||
const streamEvent: StdoutMessage = {
|
||||
type: 'stream_event',
|
||||
event: {
|
||||
type: 'content_block_delta',
|
||||
index: 0,
|
||||
delta: { type: 'text_delta', text: 'hello' },
|
||||
},
|
||||
}
|
||||
|
||||
await transport.write(streamEvent)
|
||||
await transport.close()
|
||||
|
||||
expect(posts).toEqual([
|
||||
{
|
||||
url: 'https://example.com/v2/session_ingress/session/session-1/events',
|
||||
data: { events: [streamEvent] },
|
||||
},
|
||||
])
|
||||
})
|
||||
|
||||
test('uses a close grace period when the final upload stalls', async () => {
|
||||
jest.useFakeTimers()
|
||||
try {
|
||||
postImpl = async () => new Promise(() => {})
|
||||
const transport = await createTransport({ closeGraceMs: 1 })
|
||||
|
||||
const writePromise = transport.write({
|
||||
type: 'result',
|
||||
subtype: 'success',
|
||||
duration_ms: 0,
|
||||
duration_api_ms: 0,
|
||||
is_error: false,
|
||||
result: 'ok',
|
||||
session_id: 'session-1',
|
||||
})
|
||||
const closePromise = transport.close().then(() => 'closed' as const)
|
||||
|
||||
expect(await settledValue(closePromise)).toBe('pending')
|
||||
|
||||
jest.advanceTimersByTime(1)
|
||||
await expect(closePromise).resolves.toBe('closed')
|
||||
await expect(writePromise).resolves.toBeUndefined()
|
||||
} finally {
|
||||
jest.useRealTimers()
|
||||
}
|
||||
})
|
||||
|
||||
test('clears the close grace timer when the final upload finishes first', async () => {
|
||||
jest.useFakeTimers()
|
||||
try {
|
||||
const setTimeoutSpy = jest.spyOn(globalThis, 'setTimeout')
|
||||
const clearTimeoutSpy = jest.spyOn(globalThis, 'clearTimeout')
|
||||
const transport = await createTransport({ closeGraceMs: 50 })
|
||||
|
||||
await transport.write({
|
||||
type: 'stream_event',
|
||||
event: {
|
||||
type: 'content_block_delta',
|
||||
index: 0,
|
||||
delta: { type: 'text_delta', text: 'hello' },
|
||||
},
|
||||
})
|
||||
await transport.close()
|
||||
|
||||
const closeGraceCallIndex = setTimeoutSpy.mock.calls.findIndex(
|
||||
call => call[1] === 50,
|
||||
)
|
||||
expect(closeGraceCallIndex).toBeGreaterThanOrEqual(0)
|
||||
const closeGraceTimer =
|
||||
setTimeoutSpy.mock.results[closeGraceCallIndex]?.value
|
||||
expect(
|
||||
clearTimeoutSpy.mock.calls.some(call => call[0] === closeGraceTimer),
|
||||
).toBe(true)
|
||||
} finally {
|
||||
jest.useRealTimers()
|
||||
}
|
||||
})
|
||||
|
||||
test('still drains and closes the uploader when the websocket close fails', async () => {
|
||||
const { WebSocketTransport } = await import('./WebSocketTransport.js')
|
||||
const closeError = new Error('websocket close failed')
|
||||
jest
|
||||
.spyOn(WebSocketTransport.prototype, 'close')
|
||||
.mockRejectedValueOnce(closeError)
|
||||
|
||||
const posts: Array<{ url: string; data: unknown }> = []
|
||||
postImpl = async (url, data) => {
|
||||
posts.push({ url, data })
|
||||
return { status: 200 }
|
||||
}
|
||||
const transport = await createTransport()
|
||||
const streamEvent: StdoutMessage = {
|
||||
type: 'stream_event',
|
||||
event: {
|
||||
type: 'content_block_delta',
|
||||
index: 0,
|
||||
delta: { type: 'text_delta', text: 'must still post' },
|
||||
},
|
||||
}
|
||||
|
||||
await transport.write(streamEvent)
|
||||
|
||||
await expect(transport.close()).rejects.toBe(closeError)
|
||||
expect(posts).toEqual([
|
||||
{
|
||||
url: 'https://example.com/v2/session_ingress/session/session-1/events',
|
||||
data: { events: [streamEvent] },
|
||||
},
|
||||
])
|
||||
})
|
||||
|
||||
test('preserves websocket close errors when uploader close also fails', async () => {
|
||||
const { WebSocketTransport } = await import('./WebSocketTransport.js')
|
||||
const { SerialBatchEventUploader } = await import(
|
||||
'./SerialBatchEventUploader.js'
|
||||
)
|
||||
const closeError = new Error('websocket close failed')
|
||||
const uploaderError = new Error('uploader close failed')
|
||||
jest
|
||||
.spyOn(WebSocketTransport.prototype, 'close')
|
||||
.mockRejectedValueOnce(closeError)
|
||||
jest
|
||||
.spyOn(SerialBatchEventUploader.prototype, 'close')
|
||||
.mockImplementationOnce(() => {
|
||||
throw uploaderError
|
||||
})
|
||||
|
||||
const transport = await createTransport()
|
||||
|
||||
await expect(transport.close()).rejects.toBe(closeError)
|
||||
})
|
||||
|
||||
test('surfaces uploader close errors when websocket close succeeds', async () => {
|
||||
const { SerialBatchEventUploader } = await import(
|
||||
'./SerialBatchEventUploader.js'
|
||||
)
|
||||
const uploaderError = new Error('uploader close failed')
|
||||
jest
|
||||
.spyOn(SerialBatchEventUploader.prototype, 'close')
|
||||
.mockImplementationOnce(() => {
|
||||
throw uploaderError
|
||||
})
|
||||
|
||||
const transport = await createTransport()
|
||||
|
||||
await expect(transport.close()).rejects.toBe(uploaderError)
|
||||
})
|
||||
})
|
||||
|
||||
async function createTransport(options?: { closeGraceMs?: number }) {
|
||||
const { HybridTransport } = await import('./HybridTransport.js')
|
||||
return new HybridTransport(
|
||||
new URL('wss://example.com/v2/session_ingress/ws/session-1'),
|
||||
{},
|
||||
'session-1',
|
||||
undefined,
|
||||
options,
|
||||
)
|
||||
}
|
||||
|
||||
async function settledValue<T>(promise: Promise<T>): Promise<T | 'pending'> {
|
||||
const pending = Symbol('pending')
|
||||
const result = await Promise.race([promise, Promise.resolve(pending)])
|
||||
return result === pending ? 'pending' : result
|
||||
}
|
||||
@@ -2,6 +2,7 @@ import axios, { type AxiosError } from 'axios'
|
||||
import type { StdoutMessage } from 'src/entrypoints/sdk/controlTypes.js'
|
||||
import { logForDebugging } from '../../utils/debug.js'
|
||||
import { logForDiagnosticsNoPII } from '../../utils/diagLogs.js'
|
||||
import { errorMessage } from '../../utils/errors.js'
|
||||
import { getSessionIngressAuthToken } from '../../utils/sessionIngressAuth.js'
|
||||
import { SerialBatchEventUploader } from './SerialBatchEventUploader.js'
|
||||
import {
|
||||
@@ -54,6 +55,7 @@ const CLOSE_GRACE_MS = 3000
|
||||
export class HybridTransport extends WebSocketTransport {
|
||||
private postUrl: string
|
||||
private uploader: SerialBatchEventUploader<StdoutMessage>
|
||||
private closeGraceMs: number
|
||||
|
||||
// stream_event delay buffer — accumulates content deltas for up to
|
||||
// BATCH_FLUSH_INTERVAL_MS before enqueueing (reduces POST count)
|
||||
@@ -68,11 +70,14 @@ export class HybridTransport extends WebSocketTransport {
|
||||
options?: WebSocketTransportOptions & {
|
||||
maxConsecutiveFailures?: number
|
||||
onBatchDropped?: (batchSize: number, failures: number) => void
|
||||
closeGraceMs?: number
|
||||
},
|
||||
) {
|
||||
super(url, headers, sessionId, refreshHeaders, options)
|
||||
const { maxConsecutiveFailures, onBatchDropped } = options ?? {}
|
||||
const { maxConsecutiveFailures, onBatchDropped, closeGraceMs } =
|
||||
options ?? {}
|
||||
this.postUrl = convertWsUrlToPostUrl(url)
|
||||
this.closeGraceMs = closeGraceMs ?? CLOSE_GRACE_MS
|
||||
this.uploader = new SerialBatchEventUploader<StdoutMessage>({
|
||||
// Large cap — session-ingress accepts arbitrary batch sizes. Events
|
||||
// naturally batch during in-flight POSTs; this just bounds the payload.
|
||||
@@ -168,30 +173,57 @@ export class HybridTransport extends WebSocketTransport {
|
||||
void this.uploader.enqueue(this.takeStreamEvents())
|
||||
}
|
||||
|
||||
override close(): void {
|
||||
if (this.streamEventTimer) {
|
||||
clearTimeout(this.streamEventTimer)
|
||||
this.streamEventTimer = null
|
||||
override async close(): Promise<void> {
|
||||
const pendingStreamEvents = this.takeStreamEvents()
|
||||
let closeError: unknown
|
||||
let didCloseThrow = false
|
||||
try {
|
||||
await super.close()
|
||||
} catch (error) {
|
||||
closeError = error
|
||||
didCloseThrow = true
|
||||
}
|
||||
|
||||
const { uploader } = this
|
||||
let uploaderCloseError: unknown
|
||||
if (uploader) {
|
||||
let closeGraceTimer: ReturnType<typeof setTimeout> | null = null
|
||||
try {
|
||||
await Promise.race([
|
||||
(async () => {
|
||||
if (pendingStreamEvents.length > 0) {
|
||||
await uploader.enqueue(pendingStreamEvents)
|
||||
}
|
||||
await uploader.flush()
|
||||
})(),
|
||||
new Promise<void>(resolve => {
|
||||
closeGraceTimer = setTimeout(resolve, this.closeGraceMs)
|
||||
}),
|
||||
])
|
||||
} catch {
|
||||
// Ignore flush errors on shutdown
|
||||
} finally {
|
||||
if (closeGraceTimer) {
|
||||
clearTimeout(closeGraceTimer)
|
||||
}
|
||||
try {
|
||||
uploader.close()
|
||||
} catch (error) {
|
||||
uploaderCloseError = error
|
||||
logForDebugging(
|
||||
`HybridTransport: uploader close failed: ${errorMessage(error)}`,
|
||||
)
|
||||
logForDiagnosticsNoPII('warn', 'cli_hybrid_uploader_close_error')
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
if (didCloseThrow) {
|
||||
throw closeError
|
||||
}
|
||||
if (uploaderCloseError) {
|
||||
throw uploaderCloseError
|
||||
}
|
||||
this.streamEventBuffer = []
|
||||
// Grace period for queued writes — fallback. replBridge teardown now
|
||||
// awaits archive between write and close (see CLOSE_GRACE_MS), so
|
||||
// archive latency is the primary drain window and this is a last
|
||||
// resort. Keep close() sync (returns immediately) but defer
|
||||
// uploader.close() so any remaining queue gets a chance to finish.
|
||||
const uploader = this.uploader
|
||||
let graceTimer: ReturnType<typeof setTimeout> | undefined
|
||||
void Promise.race([
|
||||
uploader.flush(),
|
||||
new Promise<void>(r => {
|
||||
// eslint-disable-next-line no-restricted-syntax -- need timer ref for clearTimeout
|
||||
graceTimer = setTimeout(r, CLOSE_GRACE_MS)
|
||||
}),
|
||||
]).finally(() => {
|
||||
clearTimeout(graceTimer)
|
||||
uploader.close()
|
||||
})
|
||||
super.close()
|
||||
}
|
||||
|
||||
/**
|
||||
|
||||
@@ -676,7 +676,7 @@ export class SSETransport implements Transport {
|
||||
this.onEventCallback = callback
|
||||
}
|
||||
|
||||
close(): void {
|
||||
async close(): Promise<void> {
|
||||
if (this.reconnectTimer) {
|
||||
clearTimeout(this.reconnectTimer)
|
||||
this.reconnectTimer = null
|
||||
|
||||
@@ -17,13 +17,13 @@ export interface Transport {
|
||||
* reconnection/backoff; the returned promise resolves when the initial
|
||||
* connection attempt completes (or the transport gives up).
|
||||
*/
|
||||
connect(): Promise<void>
|
||||
connect(): Promise<void> | void
|
||||
|
||||
/** Send a message to the server. */
|
||||
write(message: StdoutMessage): Promise<void>
|
||||
write(message: StdoutMessage): Promise<void> | void
|
||||
|
||||
/** Permanently close the transport. */
|
||||
close(): void
|
||||
/** Permanently close the transport, allowing async flush/cleanup. */
|
||||
close(): Promise<void> | void
|
||||
|
||||
/** Register the callback invoked with newline-delimited JSON payloads. */
|
||||
setOnData(callback: (data: string) => void): void
|
||||
@@ -34,9 +34,18 @@ export interface Transport {
|
||||
*/
|
||||
setOnClose(callback: (closeCode?: number) => void): void
|
||||
|
||||
/** Register the callback invoked when the transport connects. */
|
||||
setOnConnect?(callback: () => void): void
|
||||
|
||||
/** Register the callback invoked for transport-specific events. */
|
||||
setOnEvent?(callback: (event: unknown) => void): void
|
||||
|
||||
/** Whether the transport is currently connected. */
|
||||
isConnectedStatus(): boolean
|
||||
isConnectedStatus?(): boolean
|
||||
|
||||
/** Whether the transport has permanently closed. */
|
||||
isClosedStatus(): boolean
|
||||
isClosedStatus?(): boolean
|
||||
|
||||
/** Human-readable transport state for diagnostics. */
|
||||
getStateLabel?(): string
|
||||
}
|
||||
|
||||
@@ -553,7 +553,7 @@ export class WebSocketTransport implements Transport {
|
||||
}
|
||||
}
|
||||
|
||||
close(): void {
|
||||
async close(): Promise<void> {
|
||||
// Clear any pending reconnection timer
|
||||
if (this.reconnectTimer) {
|
||||
clearTimeout(this.reconnectTimer)
|
||||
|
||||
@@ -0,0 +1,318 @@
|
||||
import { describe, expect, test } from 'bun:test'
|
||||
import type { SDKPartialAssistantMessage } from 'src/entrypoints/sdk/controlTypes.js'
|
||||
import {
|
||||
accumulateStreamEvents,
|
||||
createStreamAccumulator,
|
||||
} from './ccrClient.js'
|
||||
|
||||
describe('accumulateStreamEvents', () => {
|
||||
test('coalesces text deltas into one full-so-far snapshot per content block', () => {
|
||||
const state = createStreamAccumulator()
|
||||
const messageStart = streamMessage('start', {
|
||||
type: 'message_start',
|
||||
message: { id: 'msg_1' },
|
||||
})
|
||||
const firstDelta = streamMessage('delta-1', {
|
||||
type: 'content_block_delta',
|
||||
index: 0,
|
||||
delta: { type: 'text_delta', text: 'hel' },
|
||||
})
|
||||
const secondDelta = streamMessage('delta-2', {
|
||||
type: 'content_block_delta',
|
||||
index: 0,
|
||||
delta: { type: 'text_delta', text: 'lo' },
|
||||
})
|
||||
|
||||
expect(
|
||||
accumulateStreamEvents([messageStart, firstDelta, secondDelta], state),
|
||||
).toEqual([
|
||||
messageStart,
|
||||
{
|
||||
type: 'stream_event',
|
||||
uuid: 'delta-1',
|
||||
session_id: 'session-1',
|
||||
parent_tool_use_id: null,
|
||||
event: {
|
||||
type: 'content_block_delta',
|
||||
index: 0,
|
||||
delta: { type: 'text_delta', text: 'hello' },
|
||||
},
|
||||
},
|
||||
])
|
||||
})
|
||||
|
||||
test('coalesces interleaved text deltas independently by content block index', () => {
|
||||
const state = createStreamAccumulator()
|
||||
const messageStart = streamMessage('start', {
|
||||
type: 'message_start',
|
||||
message: { id: 'msg_1' },
|
||||
})
|
||||
|
||||
const result = accumulateStreamEvents(
|
||||
[
|
||||
messageStart,
|
||||
streamMessage('block-0-a', {
|
||||
type: 'content_block_delta',
|
||||
index: 0,
|
||||
delta: { type: 'text_delta', text: 'he' },
|
||||
}),
|
||||
streamMessage('block-1-a', {
|
||||
type: 'content_block_delta',
|
||||
index: 1,
|
||||
delta: { type: 'text_delta', text: 'wo' },
|
||||
}),
|
||||
streamMessage('block-0-b', {
|
||||
type: 'content_block_delta',
|
||||
index: 0,
|
||||
delta: { type: 'text_delta', text: 'llo' },
|
||||
}),
|
||||
streamMessage('block-1-b', {
|
||||
type: 'content_block_delta',
|
||||
index: 1,
|
||||
delta: { type: 'text_delta', text: 'rld' },
|
||||
}),
|
||||
],
|
||||
state,
|
||||
)
|
||||
|
||||
expect(result).toHaveLength(3)
|
||||
expect(result).toEqual([
|
||||
messageStart,
|
||||
{
|
||||
type: 'stream_event',
|
||||
uuid: 'block-0-a',
|
||||
session_id: 'session-1',
|
||||
parent_tool_use_id: null,
|
||||
event: {
|
||||
type: 'content_block_delta',
|
||||
index: 0,
|
||||
delta: { type: 'text_delta', text: 'hello' },
|
||||
},
|
||||
},
|
||||
{
|
||||
type: 'stream_event',
|
||||
uuid: 'block-1-a',
|
||||
session_id: 'session-1',
|
||||
parent_tool_use_id: null,
|
||||
event: {
|
||||
type: 'content_block_delta',
|
||||
index: 1,
|
||||
delta: { type: 'text_delta', text: 'world' },
|
||||
},
|
||||
},
|
||||
])
|
||||
})
|
||||
|
||||
test('isolates text delta accumulation by session and parent tool scope', () => {
|
||||
const state = createStreamAccumulator()
|
||||
const rootStart = streamMessage('root-start', {
|
||||
type: 'message_start',
|
||||
message: { id: 'msg_root' },
|
||||
})
|
||||
const parentStart = streamMessage(
|
||||
'parent-start',
|
||||
{
|
||||
type: 'message_start',
|
||||
message: { id: 'msg_parent' },
|
||||
},
|
||||
{ parent_tool_use_id: 'tool-1' },
|
||||
)
|
||||
const otherSessionStart = streamMessage(
|
||||
'other-session-start',
|
||||
{
|
||||
type: 'message_start',
|
||||
message: { id: 'msg_other_session' },
|
||||
},
|
||||
{ session_id: 'session-2' },
|
||||
)
|
||||
|
||||
const result = accumulateStreamEvents(
|
||||
[
|
||||
rootStart,
|
||||
streamMessage('root-a', {
|
||||
type: 'content_block_delta',
|
||||
index: 0,
|
||||
delta: { type: 'text_delta', text: 'root-' },
|
||||
}),
|
||||
parentStart,
|
||||
streamMessage(
|
||||
'parent-a',
|
||||
{
|
||||
type: 'content_block_delta',
|
||||
index: 0,
|
||||
delta: { type: 'text_delta', text: 'parent-' },
|
||||
},
|
||||
{ parent_tool_use_id: 'tool-1' },
|
||||
),
|
||||
otherSessionStart,
|
||||
streamMessage(
|
||||
'other-session-a',
|
||||
{
|
||||
type: 'content_block_delta',
|
||||
index: 0,
|
||||
delta: { type: 'text_delta', text: 'other-' },
|
||||
},
|
||||
{ session_id: 'session-2' },
|
||||
),
|
||||
streamMessage('root-b', {
|
||||
type: 'content_block_delta',
|
||||
index: 0,
|
||||
delta: { type: 'text_delta', text: 'scope' },
|
||||
}),
|
||||
streamMessage(
|
||||
'parent-b',
|
||||
{
|
||||
type: 'content_block_delta',
|
||||
index: 0,
|
||||
delta: { type: 'text_delta', text: 'scope' },
|
||||
},
|
||||
{ parent_tool_use_id: 'tool-1' },
|
||||
),
|
||||
streamMessage(
|
||||
'other-session-b',
|
||||
{
|
||||
type: 'content_block_delta',
|
||||
index: 0,
|
||||
delta: { type: 'text_delta', text: 'scope' },
|
||||
},
|
||||
{ session_id: 'session-2' },
|
||||
),
|
||||
],
|
||||
state,
|
||||
)
|
||||
|
||||
expect(result).toEqual([
|
||||
rootStart,
|
||||
{
|
||||
type: 'stream_event',
|
||||
uuid: 'root-a',
|
||||
session_id: 'session-1',
|
||||
parent_tool_use_id: null,
|
||||
event: {
|
||||
type: 'content_block_delta',
|
||||
index: 0,
|
||||
delta: { type: 'text_delta', text: 'root-scope' },
|
||||
},
|
||||
},
|
||||
parentStart,
|
||||
{
|
||||
type: 'stream_event',
|
||||
uuid: 'parent-a',
|
||||
session_id: 'session-1',
|
||||
parent_tool_use_id: 'tool-1',
|
||||
event: {
|
||||
type: 'content_block_delta',
|
||||
index: 0,
|
||||
delta: { type: 'text_delta', text: 'parent-scope' },
|
||||
},
|
||||
},
|
||||
otherSessionStart,
|
||||
{
|
||||
type: 'stream_event',
|
||||
uuid: 'other-session-a',
|
||||
session_id: 'session-2',
|
||||
parent_tool_use_id: null,
|
||||
event: {
|
||||
type: 'content_block_delta',
|
||||
index: 0,
|
||||
delta: { type: 'text_delta', text: 'other-scope' },
|
||||
},
|
||||
},
|
||||
])
|
||||
})
|
||||
|
||||
test('resets scoped accumulation when a new message starts', () => {
|
||||
const state = createStreamAccumulator()
|
||||
const firstStart = streamMessage('first-start', {
|
||||
type: 'message_start',
|
||||
message: { id: 'msg_first' },
|
||||
})
|
||||
const secondStart = streamMessage('second-start', {
|
||||
type: 'message_start',
|
||||
message: { id: 'msg_second' },
|
||||
})
|
||||
|
||||
const result = accumulateStreamEvents(
|
||||
[
|
||||
firstStart,
|
||||
streamMessage('first-delta', {
|
||||
type: 'content_block_delta',
|
||||
index: 0,
|
||||
delta: { type: 'text_delta', text: 'first' },
|
||||
}),
|
||||
secondStart,
|
||||
streamMessage('second-delta', {
|
||||
type: 'content_block_delta',
|
||||
index: 0,
|
||||
delta: { type: 'text_delta', text: 'second' },
|
||||
}),
|
||||
],
|
||||
state,
|
||||
)
|
||||
|
||||
expect(result).toEqual([
|
||||
firstStart,
|
||||
{
|
||||
type: 'stream_event',
|
||||
uuid: 'first-delta',
|
||||
session_id: 'session-1',
|
||||
parent_tool_use_id: null,
|
||||
event: {
|
||||
type: 'content_block_delta',
|
||||
index: 0,
|
||||
delta: { type: 'text_delta', text: 'first' },
|
||||
},
|
||||
},
|
||||
secondStart,
|
||||
{
|
||||
type: 'stream_event',
|
||||
uuid: 'second-delta',
|
||||
session_id: 'session-1',
|
||||
parent_tool_use_id: null,
|
||||
event: {
|
||||
type: 'content_block_delta',
|
||||
index: 0,
|
||||
delta: { type: 'text_delta', text: 'second' },
|
||||
},
|
||||
},
|
||||
])
|
||||
expect(state.byMessage.has('msg_first')).toBe(false)
|
||||
expect(state.byMessage.has('msg_second')).toBe(true)
|
||||
})
|
||||
|
||||
test('passes malformed or non-text content deltas through unchanged', () => {
|
||||
const state = createStreamAccumulator()
|
||||
const malformedDelta = streamMessage('malformed-delta', {
|
||||
type: 'content_block_delta',
|
||||
index: 0,
|
||||
delta: null,
|
||||
})
|
||||
const imageDelta = streamMessage('image-delta', {
|
||||
type: 'content_block_delta',
|
||||
index: 1,
|
||||
delta: { type: 'input_json_delta', partial_json: '{}' },
|
||||
})
|
||||
|
||||
expect(accumulateStreamEvents([malformedDelta, imageDelta], state)).toEqual([
|
||||
malformedDelta,
|
||||
imageDelta,
|
||||
])
|
||||
})
|
||||
})
|
||||
|
||||
function streamMessage(
|
||||
uuid: string,
|
||||
event: Record<string, unknown>,
|
||||
overrides: Partial<
|
||||
Pick<SDKPartialAssistantMessage, 'parent_tool_use_id' | 'session_id'>
|
||||
> = {},
|
||||
): SDKPartialAssistantMessage {
|
||||
return {
|
||||
type: 'stream_event',
|
||||
event,
|
||||
parent_tool_use_id: null,
|
||||
uuid,
|
||||
session_id: 'session-1',
|
||||
...overrides,
|
||||
}
|
||||
}
|
||||
@@ -1,5 +1,4 @@
|
||||
import { randomUUID } from 'crypto'
|
||||
import type { RawMessageStreamEvent } from '@anthropic-ai/sdk/resources/messages.mjs'
|
||||
import type {
|
||||
SDKPartialAssistantMessage,
|
||||
StdoutMessage,
|
||||
@@ -96,6 +95,17 @@ type CoalescedStreamEvent = {
|
||||
}
|
||||
}
|
||||
|
||||
type MessageStartEventPayload = EventPayload & {
|
||||
type: 'message_start'
|
||||
message: { id: string }
|
||||
}
|
||||
|
||||
type TextDeltaEventPayload = EventPayload & {
|
||||
type: 'content_block_delta'
|
||||
index: number
|
||||
delta: { type: 'text_delta'; text: string }
|
||||
}
|
||||
|
||||
/**
|
||||
* Accumulator state for text_delta coalescing. Keyed by API message ID so
|
||||
* lifetime is tied to the assistant message — cleared when the complete
|
||||
@@ -125,6 +135,36 @@ function scopeKey(m: {
|
||||
return `${m.session_id}:${m.parent_tool_use_id ?? ''}`
|
||||
}
|
||||
|
||||
function isRecord(value: unknown): value is Record<string, unknown> {
|
||||
return typeof value === 'object' && value !== null
|
||||
}
|
||||
|
||||
function toEventPayload(event: Record<string, unknown>): EventPayload | null {
|
||||
return typeof event.type === 'string' ? (event as EventPayload) : null
|
||||
}
|
||||
|
||||
function isMessageStartEventPayload(
|
||||
event: EventPayload,
|
||||
): event is MessageStartEventPayload {
|
||||
return (
|
||||
event.type === 'message_start' &&
|
||||
isRecord(event.message) &&
|
||||
typeof event.message.id === 'string'
|
||||
)
|
||||
}
|
||||
|
||||
function isTextDeltaEventPayload(
|
||||
event: EventPayload,
|
||||
): event is TextDeltaEventPayload {
|
||||
return (
|
||||
event.type === 'content_block_delta' &&
|
||||
typeof event.index === 'number' &&
|
||||
isRecord(event.delta) &&
|
||||
event.delta.type === 'text_delta' &&
|
||||
typeof event.delta.text === 'string'
|
||||
)
|
||||
}
|
||||
|
||||
/**
|
||||
* Accumulate text_delta stream_events into full-so-far snapshots per content
|
||||
* block. Each flush emits ONE event per touched block containing the FULL
|
||||
@@ -149,12 +189,23 @@ export function accumulateStreamEvents(
|
||||
// rewrite the same entry instead of emitting one event per delta.
|
||||
const touched = new Map<string[], CoalescedStreamEvent>()
|
||||
for (const msg of buffer) {
|
||||
// RawMessageStreamEventPlaceholder is z.unknown() in the SDK schema
|
||||
// snapshot, so msg.event types as unknown — restore the concrete API
|
||||
// stream-event union for narrowing.
|
||||
const event = msg.event as RawMessageStreamEvent
|
||||
if (!isRecord(msg.event)) {
|
||||
out.push(msg)
|
||||
continue
|
||||
}
|
||||
|
||||
const event = toEventPayload(msg.event)
|
||||
if (!event) {
|
||||
out.push(msg)
|
||||
continue
|
||||
}
|
||||
|
||||
switch (event.type) {
|
||||
case 'message_start': {
|
||||
if (!isMessageStartEventPayload(event)) {
|
||||
out.push(msg)
|
||||
break
|
||||
}
|
||||
const id = event.message.id
|
||||
const prevId = state.scopeToMessage.get(scopeKey(msg))
|
||||
if (prevId) state.byMessage.delete(prevId)
|
||||
@@ -164,7 +215,7 @@ export function accumulateStreamEvents(
|
||||
break
|
||||
}
|
||||
case 'content_block_delta': {
|
||||
if (event.delta.type !== 'text_delta') {
|
||||
if (!isTextDeltaEventPayload(event)) {
|
||||
out.push(msg)
|
||||
break
|
||||
}
|
||||
@@ -984,7 +1035,7 @@ export class CCRClient {
|
||||
}
|
||||
|
||||
/** Clean up uploaders and timers. */
|
||||
close(): void {
|
||||
async close(): Promise<void> {
|
||||
this.closed = true
|
||||
this.stopHeartbeat()
|
||||
unregisterSessionActivityCallback()
|
||||
|
||||
@@ -3,6 +3,7 @@ import { mkdtempSync, rmSync } from 'node:fs'
|
||||
import { tmpdir } from 'node:os'
|
||||
import { join } from 'node:path'
|
||||
import {
|
||||
getAutoCompactThreshold,
|
||||
MAX_CONSECUTIVE_AUTOCOMPACT_FAILURES,
|
||||
type AutoCompactTrackingState,
|
||||
} from '../services/compact/autoCompact.js'
|
||||
@@ -24,6 +25,7 @@ const SAVED_ENV = {
|
||||
DISABLE_AUTO_COMPACT: process.env.DISABLE_AUTO_COMPACT,
|
||||
DISABLE_COMPACT: process.env.DISABLE_COMPACT,
|
||||
}
|
||||
|
||||
let savedAutoCompactEnabled: boolean | undefined
|
||||
let tempDir: string | undefined
|
||||
|
||||
@@ -75,26 +77,9 @@ function userMessage(content: string): Message {
|
||||
}
|
||||
}
|
||||
|
||||
function highContextMessages(): Message[] {
|
||||
return [
|
||||
{
|
||||
type: 'assistant',
|
||||
message: {
|
||||
id: 'msg-high-context',
|
||||
role: 'assistant',
|
||||
content: [{ type: 'text', text: 'previous response' }],
|
||||
usage: {
|
||||
input_tokens: 170_000,
|
||||
output_tokens: 1_000,
|
||||
cache_creation_input_tokens: 0,
|
||||
cache_read_input_tokens: 0,
|
||||
},
|
||||
},
|
||||
uuid: `assistant-${Math.random()}` as Message['uuid'],
|
||||
timestamp: new Date().toISOString(),
|
||||
} as unknown as Message,
|
||||
userMessage('continue'),
|
||||
]
|
||||
function overAutoCompactThresholdMessage(): Message {
|
||||
const threshold = getAutoCompactThreshold('claude-sonnet-4')
|
||||
return userMessage('x'.repeat((threshold + 1_000) * 4))
|
||||
}
|
||||
|
||||
function toolUseContext() {
|
||||
@@ -167,7 +152,7 @@ async function drain<T, TReturn>(
|
||||
}
|
||||
|
||||
test('active auto-compact cooldown blocks before model call with cooldown guidance', async () => {
|
||||
const messages = highContextMessages()
|
||||
const messages = [overAutoCompactThresholdMessage()]
|
||||
const nextRetryAtMs = Date.now() + 60_000
|
||||
const callModel = mock(() => {
|
||||
throw new Error('model should not be called while autocompact cools down')
|
||||
@@ -222,7 +207,7 @@ test('active auto-compact cooldown blocks before model call with cooldown guidan
|
||||
})
|
||||
|
||||
test('auto-compact cooldown tracking is carried into the next query call', async () => {
|
||||
const messages = highContextMessages()
|
||||
const messages = [overAutoCompactThresholdMessage()]
|
||||
const nextRetryAtMs = Date.now() + 60_000
|
||||
const seenTracking: Array<AutoCompactTrackingState | undefined> = []
|
||||
const callModel = mock(() => {
|
||||
@@ -335,6 +320,54 @@ test('post-compact turn tracking callback publishes a fresh object', async () =>
|
||||
expect(initialTracking.turnCounter).toBe(0)
|
||||
})
|
||||
|
||||
test('persisted breaker state does not block when auto-compact is disabled', async () => {
|
||||
process.env.DISABLE_AUTO_COMPACT = '1'
|
||||
const initialTracking: AutoCompactTrackingState = {
|
||||
compacted: false,
|
||||
turnId: 'turn',
|
||||
turnCounter: 0,
|
||||
consecutiveFailures: MAX_CONSECUTIVE_AUTOCOMPACT_FAILURES,
|
||||
nextRetryAtMs: Date.now() + 60_000,
|
||||
}
|
||||
const callModel = mock(async function* () {
|
||||
yield assistantToolUseMessage()
|
||||
})
|
||||
const deps = {
|
||||
callModel,
|
||||
microcompact: mock(async (input: Message[]) => ({
|
||||
messages: input,
|
||||
})),
|
||||
autocompact: mock(async () => ({
|
||||
wasCompacted: false,
|
||||
})),
|
||||
uuid: () => 'test-uuid',
|
||||
} as never
|
||||
|
||||
const { yielded, terminal } = await drain(
|
||||
query({
|
||||
messages: [overAutoCompactThresholdMessage()],
|
||||
systemPrompt: asSystemPrompt([]),
|
||||
userContext: {},
|
||||
systemContext: {},
|
||||
canUseTool,
|
||||
toolUseContext: toolUseContext(),
|
||||
querySource: 'repl_main_thread',
|
||||
maxTurns: 1,
|
||||
deps,
|
||||
autoCompactTracking: initialTracking,
|
||||
}),
|
||||
)
|
||||
|
||||
expect(callModel).toHaveBeenCalledTimes(1)
|
||||
expect(terminal.reason).toBe('max_turns')
|
||||
expect(
|
||||
yielded.some(
|
||||
message =>
|
||||
(message as { isApiErrorMessage?: boolean }).isApiErrorMessage === true,
|
||||
),
|
||||
).toBe(false)
|
||||
})
|
||||
|
||||
test('breaker metadata tracking callback publishes a fresh object', async () => {
|
||||
const initialTracking: AutoCompactTrackingState = {
|
||||
compacted: false,
|
||||
@@ -365,7 +398,7 @@ test('breaker metadata tracking callback publishes a fresh object', async () =>
|
||||
|
||||
const { terminal } = await drain(
|
||||
query({
|
||||
messages: highContextMessages(),
|
||||
messages: [overAutoCompactThresholdMessage()],
|
||||
systemPrompt: asSystemPrompt([]),
|
||||
userContext: {},
|
||||
systemContext: {},
|
||||
|
||||
@@ -1,5 +1,9 @@
|
||||
import { describe, expect, test } from 'bun:test'
|
||||
import { afterEach, beforeEach, describe, expect, test } from 'bun:test'
|
||||
import type { Tool } from '../Tool.js'
|
||||
import {
|
||||
acquireSharedMutationLock,
|
||||
releaseSharedMutationLock,
|
||||
} from '../test/sharedMutationLock.js'
|
||||
import { TOOL_SEARCH_TOOL_NAME } from '../tools/ToolSearchTool/constants.js'
|
||||
import { countMcpToolTokens } from './analyzeContext.js'
|
||||
import { createRequestSizeReport } from './requestSizeBreakdown.js'
|
||||
@@ -25,6 +29,12 @@ function makeToolSearchTool(): Tool {
|
||||
|
||||
const emptyPermissionContext = async () => ({ mode: 'default' }) as never
|
||||
const countToolDefinitions = async () => 1_500
|
||||
const savedToolSearchEnv = {
|
||||
ENABLE_TOOL_SEARCH: process.env.ENABLE_TOOL_SEARCH,
|
||||
CLAUDE_CODE_DISABLE_EXPERIMENTAL_BETAS:
|
||||
process.env.CLAUDE_CODE_DISABLE_EXPERIMENTAL_BETAS,
|
||||
ANTHROPIC_BASE_URL: process.env.ANTHROPIC_BASE_URL,
|
||||
}
|
||||
|
||||
function makeContextData(overrides: Partial<ContextData> = {}): ContextData {
|
||||
return {
|
||||
@@ -45,6 +55,27 @@ function makeContextData(overrides: Partial<ContextData> = {}): ContextData {
|
||||
}
|
||||
|
||||
describe('countMcpToolTokens', () => {
|
||||
beforeEach(async () => {
|
||||
await acquireSharedMutationLock('utils/analyzeContext.mcp.test.ts')
|
||||
process.env.ENABLE_TOOL_SEARCH = 'true'
|
||||
delete process.env.CLAUDE_CODE_DISABLE_EXPERIMENTAL_BETAS
|
||||
delete process.env.ANTHROPIC_BASE_URL
|
||||
})
|
||||
|
||||
afterEach(() => {
|
||||
try {
|
||||
for (const [key, value] of Object.entries(savedToolSearchEnv)) {
|
||||
if (value === undefined) {
|
||||
delete process.env[key]
|
||||
} else {
|
||||
process.env[key] = value
|
||||
}
|
||||
}
|
||||
} finally {
|
||||
releaseSharedMutationLock()
|
||||
}
|
||||
})
|
||||
|
||||
test('marks MCP tools loaded and request-size groups them by server when Tool Search is not deferred', async () => {
|
||||
const result = await countMcpToolTokens(
|
||||
[makeMcpTool('mcp__alpha__search'), makeMcpTool('mcp__beta__list')],
|
||||
|
||||
@@ -6,6 +6,17 @@ import {
|
||||
|
||||
const originalSubagentModel = process.env.CLAUDE_CODE_SUBAGENT_MODEL
|
||||
const originalOpenAIModel = process.env.OPENAI_MODEL
|
||||
const originalDefaultModelEnv = {
|
||||
ANTHROPIC_DEFAULT_OPUS_MODEL: process.env.ANTHROPIC_DEFAULT_OPUS_MODEL,
|
||||
ANTHROPIC_DEFAULT_OPUS_MODEL_SUPPORTED_CAPABILITIES:
|
||||
process.env.ANTHROPIC_DEFAULT_OPUS_MODEL_SUPPORTED_CAPABILITIES,
|
||||
ANTHROPIC_DEFAULT_SONNET_MODEL: process.env.ANTHROPIC_DEFAULT_SONNET_MODEL,
|
||||
ANTHROPIC_DEFAULT_SONNET_MODEL_SUPPORTED_CAPABILITIES:
|
||||
process.env.ANTHROPIC_DEFAULT_SONNET_MODEL_SUPPORTED_CAPABILITIES,
|
||||
ANTHROPIC_DEFAULT_HAIKU_MODEL: process.env.ANTHROPIC_DEFAULT_HAIKU_MODEL,
|
||||
ANTHROPIC_DEFAULT_HAIKU_MODEL_SUPPORTED_CAPABILITIES:
|
||||
process.env.ANTHROPIC_DEFAULT_HAIKU_MODEL_SUPPORTED_CAPABILITIES,
|
||||
}
|
||||
const allowedModelsRef: { value?: string[] } = { value: undefined }
|
||||
|
||||
type MockProvider =
|
||||
@@ -65,6 +76,10 @@ describe('getAgentModel provider-aware fallback', () => {
|
||||
beforeEach(async () => {
|
||||
await acquireSharedMutationLock('utils/model/agent.test.ts')
|
||||
delete process.env.CLAUDE_CODE_SUBAGENT_MODEL
|
||||
delete process.env.OPENAI_MODEL
|
||||
for (const key of Object.keys(originalDefaultModelEnv)) {
|
||||
delete process.env[key]
|
||||
}
|
||||
setAvailableModelsForTest()
|
||||
mockModelAllowlist()
|
||||
})
|
||||
@@ -84,6 +99,13 @@ describe('getAgentModel provider-aware fallback', () => {
|
||||
} else {
|
||||
process.env.OPENAI_MODEL = originalOpenAIModel
|
||||
}
|
||||
for (const [key, value] of Object.entries(originalDefaultModelEnv)) {
|
||||
if (value === undefined) {
|
||||
delete process.env[key]
|
||||
} else {
|
||||
process.env[key] = value
|
||||
}
|
||||
}
|
||||
} finally {
|
||||
releaseSharedMutationLock()
|
||||
}
|
||||
|
||||
@@ -60,6 +60,27 @@ const SAVED_ENV = {
|
||||
OPENAI_BASE_URL: process.env.OPENAI_BASE_URL,
|
||||
CODEX_API_KEY: process.env.CODEX_API_KEY,
|
||||
CHATGPT_ACCOUNT_ID: process.env.CHATGPT_ACCOUNT_ID,
|
||||
ANTHROPIC_DEFAULT_OPUS_MODEL: process.env.ANTHROPIC_DEFAULT_OPUS_MODEL,
|
||||
ANTHROPIC_DEFAULT_OPUS_MODEL_NAME:
|
||||
process.env.ANTHROPIC_DEFAULT_OPUS_MODEL_NAME,
|
||||
ANTHROPIC_DEFAULT_OPUS_MODEL_DESCRIPTION:
|
||||
process.env.ANTHROPIC_DEFAULT_OPUS_MODEL_DESCRIPTION,
|
||||
ANTHROPIC_DEFAULT_OPUS_MODEL_SUPPORTED_CAPABILITIES:
|
||||
process.env.ANTHROPIC_DEFAULT_OPUS_MODEL_SUPPORTED_CAPABILITIES,
|
||||
ANTHROPIC_DEFAULT_SONNET_MODEL: process.env.ANTHROPIC_DEFAULT_SONNET_MODEL,
|
||||
ANTHROPIC_DEFAULT_SONNET_MODEL_NAME:
|
||||
process.env.ANTHROPIC_DEFAULT_SONNET_MODEL_NAME,
|
||||
ANTHROPIC_DEFAULT_SONNET_MODEL_DESCRIPTION:
|
||||
process.env.ANTHROPIC_DEFAULT_SONNET_MODEL_DESCRIPTION,
|
||||
ANTHROPIC_DEFAULT_SONNET_MODEL_SUPPORTED_CAPABILITIES:
|
||||
process.env.ANTHROPIC_DEFAULT_SONNET_MODEL_SUPPORTED_CAPABILITIES,
|
||||
ANTHROPIC_DEFAULT_HAIKU_MODEL: process.env.ANTHROPIC_DEFAULT_HAIKU_MODEL,
|
||||
ANTHROPIC_DEFAULT_HAIKU_MODEL_NAME:
|
||||
process.env.ANTHROPIC_DEFAULT_HAIKU_MODEL_NAME,
|
||||
ANTHROPIC_DEFAULT_HAIKU_MODEL_DESCRIPTION:
|
||||
process.env.ANTHROPIC_DEFAULT_HAIKU_MODEL_DESCRIPTION,
|
||||
ANTHROPIC_DEFAULT_HAIKU_MODEL_SUPPORTED_CAPABILITIES:
|
||||
process.env.ANTHROPIC_DEFAULT_HAIKU_MODEL_SUPPORTED_CAPABILITIES,
|
||||
}
|
||||
// `model` is a legacy loose key not declared on GlobalConfig.
|
||||
const savedModel = (getGlobalConfig() as GlobalConfig & Record<string, unknown>).model
|
||||
@@ -94,6 +115,18 @@ beforeEach(async () => {
|
||||
delete process.env.OPENAI_BASE_URL
|
||||
delete process.env.CODEX_API_KEY
|
||||
delete process.env.CHATGPT_ACCOUNT_ID
|
||||
delete process.env.ANTHROPIC_DEFAULT_OPUS_MODEL
|
||||
delete process.env.ANTHROPIC_DEFAULT_OPUS_MODEL_NAME
|
||||
delete process.env.ANTHROPIC_DEFAULT_OPUS_MODEL_DESCRIPTION
|
||||
delete process.env.ANTHROPIC_DEFAULT_OPUS_MODEL_SUPPORTED_CAPABILITIES
|
||||
delete process.env.ANTHROPIC_DEFAULT_SONNET_MODEL
|
||||
delete process.env.ANTHROPIC_DEFAULT_SONNET_MODEL_NAME
|
||||
delete process.env.ANTHROPIC_DEFAULT_SONNET_MODEL_DESCRIPTION
|
||||
delete process.env.ANTHROPIC_DEFAULT_SONNET_MODEL_SUPPORTED_CAPABILITIES
|
||||
delete process.env.ANTHROPIC_DEFAULT_HAIKU_MODEL
|
||||
delete process.env.ANTHROPIC_DEFAULT_HAIKU_MODEL_NAME
|
||||
delete process.env.ANTHROPIC_DEFAULT_HAIKU_MODEL_DESCRIPTION
|
||||
delete process.env.ANTHROPIC_DEFAULT_HAIKU_MODEL_SUPPORTED_CAPABILITIES
|
||||
saveGlobalConfig(current => ({
|
||||
...current,
|
||||
model: undefined,
|
||||
@@ -356,4 +389,3 @@ test('default helpers do not leak claude-* names to Xiaomi MiMo', async () => {
|
||||
expect(model.toLowerCase()).not.toContain('opus')
|
||||
}
|
||||
})
|
||||
|
||||
|
||||
@@ -8,6 +8,10 @@
|
||||
"src/state/AppState.tsx",
|
||||
"src/optionalModules.d.ts",
|
||||
"src/optionalModules.types.test.ts",
|
||||
"src/bridge/replBridgeTransport.types.test.ts",
|
||||
"src/bridge/replBridgeTransport.ts",
|
||||
"src/bridge/replBridge.ts",
|
||||
"src/global.d.ts",
|
||||
"src/types/utils.types.test.ts",
|
||||
"src/types/utils.ts"
|
||||
],
|
||||
|
||||
Reference in New Issue
Block a user