fix: harden replay and pi-ai request boundaries
This commit is contained in:
@@ -6,7 +6,7 @@
|
|||||||
*/
|
*/
|
||||||
|
|
||||||
import type { Context } from 'cordis'
|
import type { Context } from 'cordis'
|
||||||
import type { FinishReason, GenerateOptions, LlmCallConfig, Message } from '@deepseek-ai/dsh-llm'
|
import type { ContentBlock, FinishReason, GenerateOptions, LlmCallConfig, Message } from '@deepseek-ai/dsh-llm'
|
||||||
import { isDeepStrictEqual } from 'node:util'
|
import { isDeepStrictEqual } from 'node:util'
|
||||||
import { BlockAssembler, HarnessError, deepFreeze } from '@deepseek-ai/dsh-llm'
|
import { BlockAssembler, HarnessError, deepFreeze } from '@deepseek-ai/dsh-llm'
|
||||||
import { agentEvents, assembleContextFor } from '@deepseek-ai/dsh-agent'
|
import { agentEvents, assembleContextFor } from '@deepseek-ai/dsh-agent'
|
||||||
@@ -529,20 +529,22 @@ async function runStep(
|
|||||||
|
|
||||||
if (assembler.finish.kind === 'max-tokens') {
|
if (assembler.finish.kind === 'max-tokens') {
|
||||||
const assembled = assembler.message()
|
const assembled = assembler.message()
|
||||||
|
const assembledContent = structuredClone(assembled.content)
|
||||||
let message: Message = withoutToolCalls(assembled)
|
let message: Message = withoutToolCalls(assembled)
|
||||||
message = withoutToolCalls(await events.waterfall('agent/step-result', turn, step, message, () => Promise.resolve(message)))
|
message = withoutToolCalls(await events.waterfall('agent/step-result', turn, step, message, () => Promise.resolve(message)))
|
||||||
// Preserve usage even when max-token truncation produced no content.
|
// Preserve usage even when max-token truncation produced no content.
|
||||||
recordAssistantMessage(session, turn, step, header.config, assembled, message, assembler, chunkSeqs)
|
recordAssistantMessage(session, turn, step, header.config, assembledContent, message, assembler, chunkSeqs)
|
||||||
return { hadToolCalls: false, finish: assembler.finish }
|
return { hadToolCalls: false, finish: assembler.finish }
|
||||||
}
|
}
|
||||||
|
|
||||||
// Record the post-waterfall message that tool dispatch uses.
|
// Record the post-waterfall message that tool dispatch uses.
|
||||||
const assembled = assembler.message()
|
const assembled = assembler.message()
|
||||||
|
const assembledContent = structuredClone(assembled.content)
|
||||||
let message: Message = assembled
|
let message: Message = assembled
|
||||||
message = await events.waterfall('agent/step-result', turn, step, message, () => Promise.resolve(message))
|
message = await events.waterfall('agent/step-result', turn, step, message, () => Promise.resolve(message))
|
||||||
|
|
||||||
// Empty messages exist only to carry usage; the helper also omits empty chunk provenance.
|
// Empty messages exist only to carry usage; the helper also omits empty chunk provenance.
|
||||||
recordAssistantMessage(session, turn, step, header.config, assembled, message, assembler, chunkSeqs)
|
recordAssistantMessage(session, turn, step, header.config, assembledContent, message, assembler, chunkSeqs)
|
||||||
|
|
||||||
// Tool execution stays sequential; recheck abort around each normalized result.
|
// Tool execution stays sequential; recheck abort around each normalized result.
|
||||||
const toolCalls = message.content.filter(block => block.type === 'tool-call')
|
const toolCalls = message.content.filter(block => block.type === 'tool-call')
|
||||||
@@ -600,7 +602,7 @@ function recordAssistantMessage(
|
|||||||
turn: number,
|
turn: number,
|
||||||
step: number,
|
step: number,
|
||||||
config: LlmCallConfig,
|
config: LlmCallConfig,
|
||||||
assembled: Message,
|
assembledContent: ContentBlock[],
|
||||||
message: Message,
|
message: Message,
|
||||||
assembler: BlockAssembler,
|
assembler: BlockAssembler,
|
||||||
chunkSeqs: number[],
|
chunkSeqs: number[],
|
||||||
@@ -615,7 +617,7 @@ function recordAssistantMessage(
|
|||||||
provenance: assistantProvenance(
|
provenance: assistantProvenance(
|
||||||
config,
|
config,
|
||||||
assembler.replayState,
|
assembler.replayState,
|
||||||
isDeepStrictEqual(message.content, assembled.content),
|
isDeepStrictEqual(message.content, assembledContent),
|
||||||
),
|
),
|
||||||
...assembler.usage === undefined ? {} : { usage: assembler.usage },
|
...assembler.usage === undefined ? {} : { usage: assembler.usage },
|
||||||
},
|
},
|
||||||
|
|||||||
@@ -110,6 +110,26 @@ describe('session log records what agent/step-result actually produced', () => {
|
|||||||
provider: 'mock', model: 'next-model', replayState,
|
provider: 'mock', model: 'next-model', replayState,
|
||||||
})
|
})
|
||||||
})
|
})
|
||||||
|
|
||||||
|
it('drops adapter replay state when step-result mutates assembled content in place', async () => {
|
||||||
|
const response = textResponse('original')
|
||||||
|
response[response.length - 1] = { type: 'finish', reason: { kind: 'stop' }, replayState: { private: 'state' } }
|
||||||
|
const adapter = new MockAdapter([response])
|
||||||
|
const ctx = await harness(adapter)
|
||||||
|
ctx.on('agent/step-result', async (_agent, _turn, _step, message) => {
|
||||||
|
const block = message.content[0]
|
||||||
|
if (block?.type === 'text') block.text = 'mutated'
|
||||||
|
return message
|
||||||
|
})
|
||||||
|
const agent = ctx.agentLoop.create(AgentId('mutated-replay-state'), { provider: 'mock', model: 'next-model' })
|
||||||
|
|
||||||
|
send(agent, 'go')
|
||||||
|
await waitForIdle(ctx, agent)
|
||||||
|
|
||||||
|
const recorded = agent.session.events.find(event => event.type === 'assistant/message')
|
||||||
|
expect(recorded?.type === 'assistant/message' && recorded.data.content).toEqual([{ type: 'text', text: 'mutated' }])
|
||||||
|
expect(recorded?.type === 'assistant/message' && recorded.data.provenance.replayState).toBeUndefined()
|
||||||
|
})
|
||||||
})
|
})
|
||||||
|
|
||||||
describe('abort during tool execution ends the turn', () => {
|
describe('abort during tool execution ends the turn', () => {
|
||||||
|
|||||||
@@ -53,6 +53,16 @@ function profileOptions(profile: PiAiProviderProfile): SimpleStreamOptions {
|
|||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
|
/** Merge deployment headers while removing case-insensitive attribution collisions. */
|
||||||
|
function requestHeaders(headers: Readonly<Record<string, string>> | undefined): Record<string, string> {
|
||||||
|
const attribution = attributionHeaders()
|
||||||
|
const reserved = new Set(Object.keys(attribution).map(name => name.toLowerCase()))
|
||||||
|
return {
|
||||||
|
...Object.fromEntries(Object.entries(headers ?? {}).filter(([name]) => !reserved.has(name.toLowerCase()))),
|
||||||
|
...attribution,
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
/**
|
/**
|
||||||
* pi-ai-backed multi-provider adapter. Model descriptors are resolved for each
|
* pi-ai-backed multi-provider adapter. Model descriptors are resolved for each
|
||||||
* request, so models need not be registered during the Cordis lifecycle.
|
* request, so models need not be registered during the Cordis lifecycle.
|
||||||
@@ -103,7 +113,7 @@ export class PiAiAdapter extends LlmAdapter {
|
|||||||
signal: controller.signal,
|
signal: controller.signal,
|
||||||
// Profile headers are deployment-owned; attribution names are
|
// Profile headers are deployment-owned; attribution names are
|
||||||
// Harness-owned and therefore win collisions.
|
// Harness-owned and therefore win collisions.
|
||||||
headers: { ...profile.headers, ...attributionHeaders() },
|
headers: requestHeaders(profile.headers),
|
||||||
})
|
})
|
||||||
yield* toStreamChunks(events)
|
yield* toStreamChunks(events)
|
||||||
} finally {
|
} finally {
|
||||||
|
|||||||
@@ -58,10 +58,10 @@ const profile = z.object({
|
|||||||
thinkingBudgets,
|
thinkingBudgets,
|
||||||
cacheRetention: z.union(['none', 'short', 'long']),
|
cacheRetention: z.union(['none', 'short', 'long']),
|
||||||
transport: z.union(['sse', 'websocket', 'websocket-cached', 'auto']),
|
transport: z.union(['sse', 'websocket', 'websocket-cached', 'auto']),
|
||||||
timeoutMs: z.number(),
|
timeoutMs: z.natural(),
|
||||||
websocketConnectTimeoutMs: z.number(),
|
websocketConnectTimeoutMs: z.natural(),
|
||||||
maxRetries: z.number(),
|
maxRetries: z.natural(),
|
||||||
maxRetryDelayMs: z.number(),
|
maxRetryDelayMs: z.natural(),
|
||||||
})
|
})
|
||||||
|
|
||||||
/** Runtime schema for {@link Config}. */
|
/** Runtime schema for {@link Config}. */
|
||||||
@@ -83,7 +83,7 @@ export function resolveProfiles(profiles: readonly PiAiProviderProfile[]): PiAiP
|
|||||||
if (source.provider.length === 0) throw new Error('llm-pi-ai: provider names must be non-empty')
|
if (source.provider.length === 0) throw new Error('llm-pi-ai: provider names must be non-empty')
|
||||||
if (!supported.has(source.provider)) throw new Error(`llm-pi-ai: unknown pi-ai provider "${source.provider}"`)
|
if (!supported.has(source.provider)) throw new Error(`llm-pi-ai: unknown pi-ai provider "${source.provider}"`)
|
||||||
if (seen.has(source.provider)) throw new Error(`llm-pi-ai: duplicate provider profile "${source.provider}"`)
|
if (seen.has(source.provider)) throw new Error(`llm-pi-ai: duplicate provider profile "${source.provider}"`)
|
||||||
if (source.apiKey !== undefined && source.apiKey.length === 0) {
|
if (source.apiKey !== undefined && source.apiKey.trim().length === 0) {
|
||||||
throw new Error(`llm-pi-ai: provider "${source.provider}" has an empty apiKey; omit it to use ambient authentication`)
|
throw new Error(`llm-pi-ai: provider "${source.provider}" has an empty apiKey; omit it to use ambient authentication`)
|
||||||
}
|
}
|
||||||
if (source.baseURL !== undefined && source.baseURL.length === 0) {
|
if (source.baseURL !== undefined && source.baseURL.length === 0) {
|
||||||
|
|||||||
@@ -90,7 +90,7 @@ describe('PiAiAdapter provider routing', () => {
|
|||||||
it('merges profile headers with Harness attribution winning', async () => {
|
it('merges profile headers with Harness attribution winning', async () => {
|
||||||
const server = await mockServer([{ events: textEvents }])
|
const server = await mockServer([{ events: textEvents }])
|
||||||
const ctx = await harness(server.url, {
|
const ctx = await harness(server.url, {
|
||||||
headers: { 'x-company': 'private', 'user-agent': 'wrong' },
|
headers: { 'x-company': 'private', 'User-Agent': 'wrong' },
|
||||||
})
|
})
|
||||||
await assemble(ctx, { model: 'deepseek-v4-flash', messages: [] })
|
await assemble(ctx, { model: 'deepseek-v4-flash', messages: [] })
|
||||||
expect(server.headers[0]?.['x-company']).toBe('private')
|
expect(server.headers[0]?.['x-company']).toBe('private')
|
||||||
@@ -219,9 +219,23 @@ describe('provider profile lifecycle', () => {
|
|||||||
expect(() => resolveProfiles([{ provider: 'not-real' }])).toThrow(/unknown/)
|
expect(() => resolveProfiles([{ provider: 'not-real' }])).toThrow(/unknown/)
|
||||||
expect(() => resolveProfiles([{ provider: 'openai' }, { provider: 'openai' }])).toThrow(/duplicate/)
|
expect(() => resolveProfiles([{ provider: 'openai' }, { provider: 'openai' }])).toThrow(/duplicate/)
|
||||||
expect(() => resolveProfiles([{ provider: 'openai', apiKey: '' }])).toThrow(/empty apiKey/)
|
expect(() => resolveProfiles([{ provider: 'openai', apiKey: '' }])).toThrow(/empty apiKey/)
|
||||||
|
expect(() => resolveProfiles([{ provider: 'openai', apiKey: ' ' }])).toThrow(/empty apiKey/)
|
||||||
expect(() => resolveProfiles([{ provider: 'openai', baseURL: '' }])).toThrow(/empty baseURL/)
|
expect(() => resolveProfiles([{ provider: 'openai', baseURL: '' }])).toThrow(/empty baseURL/)
|
||||||
})
|
})
|
||||||
|
|
||||||
|
it('rejects negative or fractional stream tunables at schema validation', () => {
|
||||||
|
const invalid = [
|
||||||
|
{ timeoutMs: -1 },
|
||||||
|
{ websocketConnectTimeoutMs: -1 },
|
||||||
|
{ maxRetries: -1 },
|
||||||
|
{ maxRetries: 0.5 },
|
||||||
|
{ maxRetryDelayMs: -1 },
|
||||||
|
]
|
||||||
|
for (const entry of invalid) {
|
||||||
|
expect(() => new LlmPiAi.Config({ providers: [{ provider: 'openai', ...entry }] })).toThrow()
|
||||||
|
}
|
||||||
|
})
|
||||||
|
|
||||||
it('constructs the adapter directly and rejects routes it does not own', async () => {
|
it('constructs the adapter directly and rejects routes it does not own', async () => {
|
||||||
const adapter = new PiAiAdapter({ profiles: [{ provider: 'openai' }] })
|
const adapter = new PiAiAdapter({ profiles: [{ provider: 'openai' }] })
|
||||||
await expect(adapter.listModels('anthropic')).rejects.toMatchObject({ code: 'NO_ADAPTER' })
|
await expect(adapter.listModels('anthropic')).rejects.toMatchObject({ code: 'NO_ADAPTER' })
|
||||||
|
|||||||
Reference in New Issue
Block a user