From 6fa2d046b394b6ed6555399bcb77359a83665383 Mon Sep 17 00:00:00 2001 From: suguanYang Date: Sun, 5 Jul 2026 23:52:01 +0800 Subject: [PATCH] Fix source selection, archive idempotency, and tool metadata --- src/agent-harness/runtime.test.ts | 75 ++++++++++++ src/agent-harness/runtime.ts | 67 ++++++++++- src/components/workspace-source-state.test.ts | 47 ++++++++ src/components/workspace-source-state.ts | 6 +- src/domains/sources/route-archive.ts | 50 ++++++++ src/domains/sources/route-service.test.ts | 112 ++++++++++++++++++ 6 files changed, 354 insertions(+), 3 deletions(-) diff --git a/src/agent-harness/runtime.test.ts b/src/agent-harness/runtime.test.ts index 5657e5c..5470690 100644 --- a/src/agent-harness/runtime.test.ts +++ b/src/agent-harness/runtime.test.ts @@ -5,6 +5,7 @@ import { buildHarnessMessages, buildHarnessSystemPrompt, createHarnessTools, + sanitizeHarnessModelMessagesForStep, } from "./runtime" import { createEvidenceLedger } from "./ledger" import type { @@ -251,6 +252,80 @@ describe("agent harness runtime", () => { expect(JSON.stringify(messages)).toContain("id=turn_1 role=assistant") expect(JSON.stringify(messages)).not.toContain("searchSources.query") }) + + it("removes provider metadata from tool-result parts while preserving tool-call metadata", () => { + const messages = sanitizeHarnessModelMessagesForStep([ + { + role: "assistant", + content: [ + { + type: "tool-call", + toolCallId: "call_1", + toolName: "retrieve", + input: { query: "q4" }, + providerOptions: { + google: { + thoughtSignature: "signature-1", + }, + }, + }, + ], + }, + { + role: "tool", + content: [ + { + type: "tool-result", + toolCallId: "call_1", + toolName: "retrieve", + output: { + type: "json", + value: { + ok: true, + }, + }, + providerOptions: { + google: { + thoughtSignature: "signature-1", + }, + }, + }, + ], + }, + ]) + + expect(messages).toEqual([ + { + role: "assistant", + content: [ + expect.objectContaining({ + type: "tool-call", + providerOptions: { + google: { + thoughtSignature: "signature-1", + }, + }, + }), + ], + }, + { + role: "tool", + content: [ + { + type: "tool-result", + toolCallId: "call_1", + toolName: "retrieve", + output: { + type: "json", + value: { + ok: true, + }, + }, + }, + ], + }, + ]) + }) }) function executeTool(tool: unknown, input: unknown): Promise { diff --git a/src/agent-harness/runtime.ts b/src/agent-harness/runtime.ts index c03baac..895dfc5 100644 --- a/src/agent-harness/runtime.ts +++ b/src/agent-harness/runtime.ts @@ -1,4 +1,10 @@ -import { stepCountIs, ToolLoopAgent, tool, type ModelMessage } from "ai" +import { + stepCountIs, + ToolLoopAgent, + tool, + type ModelMessage, + type ToolResultPart, +} from "ai" import { z } from "zod" import { createEvidenceLedger } from "./ledger" @@ -147,6 +153,9 @@ export async function runAgentHarness( model: input.model, instructions: buildHarnessSystemPrompt(input.turn), tools, + prepareStep: ({ messages: stepMessages }) => ({ + messages: sanitizeHarnessModelMessagesForStep(stepMessages), + }), stopWhen: stepCountIs(input.maxSteps ?? defaultMaxSteps), }) @@ -205,6 +214,62 @@ export async function runAgentHarness( } } +export function sanitizeHarnessModelMessagesForStep( + messages: readonly ModelMessage[], +): ModelMessage[] { + return messages.map(sanitizeHarnessModelMessageForStep) +} + +function sanitizeHarnessModelMessageForStep( + message: ModelMessage, +): ModelMessage { + if (message.role === "tool") { + return { + ...message, + providerOptions: undefined, + content: message.content.map(sanitizeToolMessageContentPart), + } + } + + if (message.role === "assistant" && Array.isArray(message.content)) { + return { + ...message, + content: message.content.map(sanitizeAssistantMessageContentPart), + } + } + + return message +} + +function sanitizeToolMessageContentPart( + part: ModelMessageForRole<"tool">["content"][number], +): ModelMessageForRole<"tool">["content"][number] { + if (part.type !== "tool-result") return part + + return removeToolResultPartProviderOptions(part) +} + +function sanitizeAssistantMessageContentPart( + part: Exclude["content"], string>[number], +): Exclude["content"], string>[number] { + if (part.type !== "tool-result") return part + + return removeToolResultPartProviderOptions(part) +} + +function removeToolResultPartProviderOptions( + part: ToolResultPart, +): ToolResultPart { + const sanitizedPart = { ...part } + delete sanitizedPart.providerOptions + return sanitizedPart +} + +type ModelMessageForRole = Extract< + ModelMessage, + { readonly role: TRole } +> + function buildRevisionFeedback(errors: readonly string[]): string { return [ "Your finalize output did not satisfy the output contract:", diff --git a/src/components/workspace-source-state.test.ts b/src/components/workspace-source-state.test.ts index 19e6cb1..bfd0273 100644 --- a/src/components/workspace-source-state.test.ts +++ b/src/components/workspace-source-state.test.ts @@ -104,6 +104,53 @@ describe("workspaceSourceState", () => { ).toBe("source_localized"); }); + it("keeps an explicit non-ready Source selected instead of falling back", () => { + const sources: readonly SourceView[] = [ + { + id: "source_pending", + title: "pending.pdf", + status: "parsing", + mimeType: "application/pdf", + excludedFromQuery: false, + }, + { + id: "source_ready", + title: "ready.pdf", + status: "ready", + mimeType: "application/pdf", + excludedFromQuery: false, + }, + ]; + + expect( + workspaceSourceState.getResolvedSelectedSourceId( + sources, + "source_pending", + ), + ).toBe("source_pending"); + }); + + it("does not resolve a stale selected Source to an unrelated ready Source", () => { + const sources: readonly SourceView[] = [ + { + id: "demo_ready", + kind: "demo", + demoSourceId: "demo_ready", + title: "demo.pdf", + status: "ready", + mimeType: "application/pdf", + excludedFromQuery: false, + }, + ]; + + expect( + workspaceSourceState.getResolvedSelectedSourceId( + sources, + "source_stale", + ), + ).toBeNull(); + }); + it("applies source query exclusions without mutating the source list", () => { const sources: readonly SourceView[] = [ { diff --git a/src/components/workspace-source-state.ts b/src/components/workspace-source-state.ts index 0373983..97c0c55 100644 --- a/src/components/workspace-source-state.ts +++ b/src/components/workspace-source-state.ts @@ -69,8 +69,10 @@ function getResolvedSelectedSourceId( sources: readonly SourceView[], selectedSourceId: string | null, ): string | null { + if (!selectedSourceId) return getFirstReadySourceId(sources) + const selectedSource = sources.find((source) => source.id === selectedSourceId) - if (selectedSource && isReadySource(selectedSource)) { + if (selectedSource) { return selectedSource.id } @@ -83,7 +85,7 @@ function getResolvedSelectedSourceId( if (localizedSource) return localizedSource.id } - return getFirstReadySourceId(sources) + return null } function isReadySource(source: SourceView): boolean { diff --git a/src/domains/sources/route-archive.ts b/src/domains/sources/route-archive.ts index 897dee3..6047256 100644 --- a/src/domains/sources/route-archive.ts +++ b/src/domains/sources/route-archive.ts @@ -72,6 +72,8 @@ const archiveSourceEffect = ( ) yield* Effect.tryPromise(() => client.documents.archive(source.knowhereDocumentId!), + ).pipe( + Effect.catchIf(isKnowhereDocumentNotFoundError, () => Effect.void), ) } @@ -92,4 +94,52 @@ const archiveSourceEffect = ( return routeResult.ok({ id: input.sourceId, archived: true as const }) }) +function isKnowhereDocumentNotFoundError(error: unknown): boolean { + return readErrorMessages(error).some(isDocumentNotFoundMessage) +} + +function readErrorMessages(error: unknown): readonly string[] { + if (error instanceof Error) { + return [ + error.message, + ...readNestedErrorMessages(error), + ].filter(isNonEmptyString) + } + + if (typeof error === "string") return [error] + + if (!isRecord(error)) return [] + + const messages: string[] = [] + if (typeof error.message === "string") messages.push(error.message) + messages.push(...readNestedErrorMessages(error)) + return messages +} + +function readNestedErrorMessages(error: unknown): readonly string[] { + if (!isRecord(error)) return [] + + const nested = [ + error.error, + error.cause, + isRecord(error.body) ? error.body.error : undefined, + ] + + return nested.flatMap((value) => + value === undefined || value === error ? [] : readErrorMessages(value), + ) +} + +function isDocumentNotFoundMessage(message: string): boolean { + return /\bdocument\s+not\s+found\b/iu.test(message) +} + +function isNonEmptyString(value: string): boolean { + return value.trim().length > 0 +} + +function isRecord(value: unknown): value is Record { + return typeof value === "object" && value !== null +} + export { createRouteArchive } diff --git a/src/domains/sources/route-service.test.ts b/src/domains/sources/route-service.test.ts index b07439c..a056427 100644 --- a/src/domains/sources/route-service.test.ts +++ b/src/domains/sources/route-service.test.ts @@ -834,6 +834,118 @@ describe("source route service", () => { expect(onUploadFinished).toHaveBeenCalledOnce(); }); + it("soft-deletes a source when Knowhere says its document is already missing", async () => { + const readySource: Source = { + ...source, + status: "ready", + knowhereJobId: null, + knowhereDocumentId: "doc_missing", + }; + const archiveDocument = vi.fn(async () => { + throw new Error("Document not found"); + }); + const softDelete = vi.fn(async () => true); + const service = createSourceRouteService({ + ensureApiKeyForWorkspace: vi.fn(async () => "jwt_123"), + ensureWorkspace: vi.fn(async () => workspace), + makeKnowhereClient: vi.fn(() => ({ + documents: { + archive: archiveDocument, + listChunks: vi.fn(async () => ({ + chunks: [], + pagination: { + page: 1, + pageSize: 1, + total: 0, + totalPages: 0, + }, + })), + }, + jobs: { + create: vi.fn(), + get: vi.fn(), + upload: vi.fn(), + }, + })), + requireUser: vi.fn(async () => ({ + id: "user_1", + email: null, + name: null, + })), + sourceService: { + findInWorkspace: vi.fn(async () => readySource), + softDelete, + }, + }); + + const result = await service.archiveSource({ + cookieHeader: "session=abc", + sourceId: "source_1", + }); + + expect(result).toEqual({ + status: 200, + body: { + id: "source_1", + archived: true, + }, + }); + expect(archiveDocument).toHaveBeenCalledWith("doc_missing"); + expect(softDelete).toHaveBeenCalledWith(workspace.id, "source_1"); + }); + + it("does not soft-delete a source when Knowhere archive fails unexpectedly", async () => { + const readySource: Source = { + ...source, + status: "ready", + knowhereJobId: null, + knowhereDocumentId: "doc_unavailable", + }; + const softDelete = vi.fn(async () => true); + const service = createSourceRouteService({ + ensureApiKeyForWorkspace: vi.fn(async () => "jwt_123"), + ensureWorkspace: vi.fn(async () => workspace), + makeKnowhereClient: vi.fn(() => ({ + documents: { + archive: vi.fn(async () => { + throw new Error("Knowhere unavailable"); + }), + listChunks: vi.fn(async () => ({ + chunks: [], + pagination: { + page: 1, + pageSize: 1, + total: 0, + totalPages: 0, + }, + })), + }, + jobs: { + create: vi.fn(), + get: vi.fn(), + upload: vi.fn(), + }, + })), + requireUser: vi.fn(async () => ({ + id: "user_1", + email: null, + name: null, + })), + sourceService: { + findInWorkspace: vi.fn(async () => readySource), + softDelete, + }, + }); + + await expect( + service.archiveSource({ + cookieHeader: "session=abc", + sourceId: "source_1", + }), + ).rejects.toThrow(); + expect(softDelete).not.toHaveBeenCalled(); + }); + it("retries a failed source from its saved original Blob", async () => { const failedSource: Source = { ...source,