From f8adef2e0ea74ef3ac072b87f5beec55233fcbc3 Mon Sep 17 00:00:00 2001 From: Makisuo Date: Thu, 6 Aug 2026 14:22:15 +0200 Subject: [PATCH 1/3] feat(chat): retry, compaction, sub-agents and permission rulesets for the chat turn MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Ports four things opencode's agent runtime has that Maple's chat turn did not, adapted where Durable Objects force a different shape than a long-lived process. `lib/llm` is untouched — it is a provider/protocol layer and deliberately owns no loop, and upstream has landed no source changes since the vendor pin. **Retry.** A stream that died mid-body killed the turn. `llm-retry.ts` classifies what is worth another attempt and `runStep` retries with backoff, bounded by a whole-turn budget so it cannot outrun `TURN_STALE_MS`. The load-bearing detail: `TransportReason.retryable` and `InvalidProviderOutputReason.retryable` are both hardcoded `false` in the vendored package, and those two are exactly how a mid-stream failure surfaces. A classifier that trusted the flag would retry only what `RequestExecutor` already retries — it would look right and do nothing. The classifier names them explicitly. Retries emit a `turn-retry` event carrying both a retraction and the progress signal, so text from an abandoned attempt is taken back rather than duplicated in a durable log. Note `Stream.groupedWithin` discards its pending buffer on upstream failure, so a provider that dies inside one batching window retracts zero. **Compaction.** `toLlmMessages` silently dropped the oldest 40 messages / 60k chars of a long investigation. It now replays a summary of the head instead. Runs post-turn — mid-turn would interleave a second model call into a stream the browser renders keyed by `messageId` — bounded by a timeout and swallowed on failure, so it can never turn a delivered answer into a failed turn. The head-drop stays as the fallback. **Sub-agents.** A `task` tool that hands a self-contained question to a read-only `explore` agent and returns only its final text, never its tool payloads. That firewall is the point: it is how the model searches broadly without burying the thread in warehouse rows. Runs in process rather than as a child Durable Object. A child DO would need a new addressing scheme surviving all four session-id parsers, a distributed abort fan-out, and either cross-DO event mirroring or a second SSE stream per task — all to solve a CPU-headroom problem the budgets already bound. In process, abort propagates for free and there is one log and one stream. **Permission rulesets.** `MUTATING_TOOL_NAMES` was a flat `Set` that could only answer "does this need approval?". Rulesets add `deny` (the model never sees the tool) and per-agent scoping. The set itself is unchanged and seeds `DEFAULT_RULESET`; a test asserts every registered tool resolves identically, so day-one behaviour is locked. Fixed in passing: `triage-agent.ts` rendered a whole Effect `Cause` — stack frames and, inside a `DatabaseError`, connection details — into the model's context and into the durable transcript. Both agents now share `mcp/tools/llm-tools.ts`. Also fixed, both found while testing sub-agents: a child's `turn-end` ended the parent's SSE read loop, in the web client and in the Durable Object's pump; and a provider emitting tool calls on the tool-less closing step would recurse forever. --- apps/api/src/chat/ChatSession.test.ts | 242 +++++++++ apps/api/src/chat/ChatSession.ts | 274 ++++++++--- apps/api/src/chat/agent.test.ts | 334 ++++++++++++- apps/api/src/chat/agent.ts | 458 ++++++++++++------ apps/api/src/chat/agents.test.ts | 99 ++++ apps/api/src/chat/agents.ts | 137 ++++++ apps/api/src/chat/context-budget.test.ts | 119 +++++ apps/api/src/chat/context-budget.ts | 137 ++++++ apps/api/src/chat/llm-retry.test.ts | 69 +++ apps/api/src/chat/llm-retry.ts | 77 +++ apps/api/src/chat/modes.ts | 30 -- apps/api/src/chat/permissions.ts | 42 ++ apps/api/src/chat/prompts.ts | 53 +- apps/api/src/chat/task-tool.test.ts | 292 +++++++++++ apps/api/src/chat/task-tool.ts | 234 +++++++++ apps/api/src/chat/turn-runner.test.ts | 121 +++++ apps/api/src/chat/turn-runner.ts | 98 +++- apps/api/src/mcp/tools/llm-tools.ts | 133 +++++ apps/api/src/mcp/tools/mutating.test.ts | 51 ++ apps/api/src/platform/Llm.test.ts | 49 +- apps/api/src/platform/Llm.ts | 91 +++- apps/api/src/workflows/triage-agent.ts | 87 +--- apps/web/src/components/ai-elements/types.ts | 18 + .../components/chat/chat-transcript.test.tsx | 51 +- .../src/components/chat/chat-transcript.tsx | 15 + apps/web/src/components/chat/task-card.tsx | 108 +++++ .../components/chat/transcript-rows.test.ts | 28 ++ apps/web/src/hooks/use-maple-chat.test.tsx | 88 +++- apps/web/src/hooks/use-maple-chat.ts | 158 +++++- packages/domain/package.json | 1 + packages/domain/src/chat-session.test.ts | 50 ++ packages/domain/src/chat-session.ts | 131 ++++- packages/domain/src/permission.test.ts | 41 ++ packages/domain/src/permission.ts | 67 +++ 34 files changed, 3609 insertions(+), 374 deletions(-) create mode 100644 apps/api/src/chat/agents.test.ts create mode 100644 apps/api/src/chat/agents.ts create mode 100644 apps/api/src/chat/context-budget.test.ts create mode 100644 apps/api/src/chat/context-budget.ts create mode 100644 apps/api/src/chat/llm-retry.test.ts create mode 100644 apps/api/src/chat/llm-retry.ts delete mode 100644 apps/api/src/chat/modes.ts create mode 100644 apps/api/src/chat/permissions.ts create mode 100644 apps/api/src/chat/task-tool.test.ts create mode 100644 apps/api/src/chat/task-tool.ts create mode 100644 apps/api/src/chat/turn-runner.test.ts create mode 100644 apps/api/src/mcp/tools/llm-tools.ts create mode 100644 apps/web/src/components/chat/task-card.tsx create mode 100644 packages/domain/src/permission.test.ts create mode 100644 packages/domain/src/permission.ts diff --git a/apps/api/src/chat/ChatSession.test.ts b/apps/api/src/chat/ChatSession.test.ts index 7cc3e791c..9af587d8d 100644 --- a/apps/api/src/chat/ChatSession.test.ts +++ b/apps/api/src/chat/ChatSession.test.ts @@ -107,6 +107,65 @@ describe("ChatSession.history", () => { assert.equal(call?.output, undefined) }) + it("retracts the text of a step that failed and was retried", () => { + const { session } = makeSession() + session.append({ type: "turn-start", messageId: "a1" }) + session.append({ type: "text-delta", messageId: "a1", text: "Hello wo" }) + session.append({ + type: "turn-retry", + messageId: "a1", + attempt: 2, + retractChars: 8, + reason: "Transport", + delayMs: 1_000, + }) + session.append({ type: "text-delta", messageId: "a1", text: "Hello world." }) + session.append({ type: "turn-end", messageId: "a1", reason: "stop" }) + + // Without the retraction the transcript reads "Hello woHello world." — permanently, since + // the fold concatenates and the log is durable. + assert.equal(session.history()[0]?.text, "Hello world.") + }) + + it("clamps an over-retraction instead of producing a negative slice", () => { + // `Stream.takeWhile(holdsTurn)` can drop appends between a delta and the retraction that + // accounts for it, so `retractChars` legitimately exceeds what landed. + const { session } = makeSession() + session.append({ type: "turn-start", messageId: "a1" }) + session.append({ type: "text-delta", messageId: "a1", text: "hi" }) + session.append({ + type: "turn-retry", + messageId: "a1", + attempt: 2, + retractChars: 999, + reason: "Transport", + delayMs: 1_000, + }) + + assert.equal(session.history()[0]?.text, "") + }) + + it("retracts only from the message that retried", () => { + const { session } = makeSession() + session.append({ type: "turn-start", messageId: "a1" }) + session.append({ type: "text-delta", messageId: "a1", text: "first turn" }) + session.append({ type: "turn-end", messageId: "a1", reason: "stop" }) + session.append({ type: "turn-start", messageId: "a2" }) + session.append({ type: "text-delta", messageId: "a2", text: "second" }) + session.append({ + type: "turn-retry", + messageId: "a2", + attempt: 2, + retractChars: 6, + reason: "Transport", + delayMs: 1_000, + }) + + const history = session.history() + assert.equal(history[0]?.text, "first turn") + assert.equal(history[1]?.text, "") + }) + it("interleaves text around a tool call within one message", () => { const { session } = makeSession() session.append({ type: "turn-start", messageId: "a1" }) @@ -362,3 +421,186 @@ describe("ChatSession.subscribe", () => { assert.equal(outcome, "still-open") }) }) + +describe("ChatSession.history — sub-agent transcripts", () => { + /** The parent announces the `task` call before the sub-agent's own events arrive. */ + const withTask = (session: ReturnType["session"], callId = "t1") => { + session.append({ type: "turn-start", messageId: "a1" }) + session.append({ + type: "tool-call", + messageId: "a1", + callId, + name: "task", + input: { subagent_type: "explore", description: "trace checkout latency" }, + }) + } + + const ref = (callId = "t1") => ({ id: callId, agent: "explore", parentMessageId: "a1" }) as const + + it("nests a sub-agent's turn under the tool call that started it", () => { + const { session } = makeSession() + withTask(session) + session.append({ type: "turn-start", messageId: "c1", task: ref() }) + session.append({ type: "text-delta", messageId: "c1", text: "p99 is 4.2s in ", task: ref() }) + session.append({ type: "text-delta", messageId: "c1", text: "checkout-api.", task: ref() }) + session.append({ + type: "tool-call", + messageId: "c1", + callId: "x1", + name: "search_traces", + input: {}, + task: ref(), + }) + session.append({ + type: "tool-result", + messageId: "c1", + callId: "x1", + output: "40k rows", + task: ref(), + }) + session.append({ type: "turn-end", messageId: "c1", reason: "stop", task: ref() }) + session.append({ type: "turn-end", messageId: "a1", reason: "stop" }) + + const history = session.history() + + // One top-level assistant message, not one per sub-agent event. + assert.lengthOf(history, 1) + const task = history[0]?.toolCalls[0]?.task + assert.equal(task?.agent, "explore") + assert.equal(task?.status, "completed") + assert.lengthOf(task?.messages ?? [], 1) + assert.equal(task?.messages[0]?.text, "p99 is 4.2s in checkout-api.") + // The child's own tool call rides in the nested transcript, not the parent's. + assert.equal(task?.messages[0]?.toolCalls[0]?.output, "40k rows") + assert.lengthOf(history[0]!.toolCalls, 1) + }) + + it("shows a sub-agent still running before its turn ends", () => { + const { session } = makeSession() + withTask(session) + session.append({ type: "turn-start", messageId: "c1", task: ref() }) + session.append({ type: "text-delta", messageId: "c1", text: "looking", task: ref() }) + + assert.equal(session.history()[0]?.toolCalls[0]?.task?.status, "running") + }) + + it("carries a failed sub-agent's status through", () => { + const { session } = makeSession() + withTask(session) + session.append({ type: "turn-start", messageId: "c1", task: ref() }) + session.append({ type: "turn-end", messageId: "c1", reason: "error", error: "boom", task: ref() }) + + assert.equal(session.history()[0]?.toolCalls[0]?.task?.status, "error") + }) + + it("drops an orphaned sub-agent event rather than corrupting the transcript", () => { + // Deny by default. A child event whose parent message or tool call is missing must not + // materialise a stray top-level assistant message — that is what would leak into both the + // browser and the next turn's replay. + const { session } = makeSession() + session.append({ type: "user-message", id: "u1", text: "hi" }) + session.append({ + type: "text-delta", + messageId: "c1", + text: "orphan", + task: { id: "nope", agent: "explore", parentMessageId: "a-missing" }, + }) + + const history = session.history() + assert.lengthOf(history, 1) + assert.equal(history[0]?.role, "user") + }) + + it("keeps two sibling sub-agents' transcripts apart", () => { + const { session } = makeSession() + session.append({ type: "turn-start", messageId: "a1" }) + for (const callId of ["t1", "t2"]) { + session.append({ type: "tool-call", messageId: "a1", callId, name: "task", input: {} }) + } + session.append({ type: "text-delta", messageId: "c1", text: "first", task: ref("t1") }) + session.append({ type: "text-delta", messageId: "c2", text: "second", task: ref("t2") }) + + const calls = session.history()[0]!.toolCalls + assert.equal(calls[0]?.task?.messages[0]?.text, "first") + assert.equal(calls[1]?.task?.messages[0]?.text, "second") + }) + + it("retracts inside a sub-agent transcript, not the parent's", () => { + const { session } = makeSession() + session.append({ type: "turn-start", messageId: "a1" }) + session.append({ type: "text-delta", messageId: "a1", text: "parent text" }) + session.append({ type: "tool-call", messageId: "a1", callId: "t1", name: "task", input: {} }) + session.append({ type: "text-delta", messageId: "c1", text: "partial", task: ref() }) + session.append({ + type: "turn-retry", + messageId: "c1", + attempt: 2, + retractChars: 7, + reason: "Transport", + delayMs: 1_000, + task: ref(), + }) + session.append({ type: "text-delta", messageId: "c1", text: "the answer", task: ref() }) + + const message = session.history()[0]! + assert.equal(message.text, "parent text") + assert.equal(message.toolCalls[0]?.task?.messages[0]?.text, "the answer") + }) +}) + +describe("ChatSession compaction", () => { + it("is inert for display — the user still sees what they actually said", () => { + // This is where Maple diverges from opencode. There, the transcript and the model input are + // the same list, so compaction reorders what the user sees. Here they are different things. + const { session } = makeSession() + session.append({ type: "user-message", id: "u1", text: "why is checkout slow?" }) + session.append({ type: "turn-start", messageId: "a1" }) + session.append({ type: "text-delta", messageId: "a1", text: "p99 is 4.2s." }) + session.append({ type: "turn-end", messageId: "a1", reason: "stop" }) + session.append({ + type: "compaction", + messageId: "a1", + summary: "checkout p99 was 4.2s", + throughSeq: 4, + }) + + const history = session.history() + assert.lengthOf(history, 2) + assert.equal(history[0]?.text, "why is checkout slow?") + assert.equal(history[1]?.text, "p99 is 4.2s.") + }) + + it("returns the most recent compaction, not the first", () => { + const { session } = makeSession() + session.append({ type: "compaction", messageId: "a1", summary: "older", throughSeq: 1 }) + session.append({ type: "user-message", id: "u1", text: "more" }) + session.append({ type: "compaction", messageId: "a2", summary: "newer", throughSeq: 2 }) + + assert.deepEqual(session.compaction(), { summary: "newer", throughSeq: 2 }) + }) + + it("reports no compaction on a fresh conversation", () => { + const { session } = makeSession() + session.append({ type: "user-message", id: "u1", text: "hi" }) + + assert.isUndefined(session.compaction()) + }) + + it("stamps every message with the seq that opened it", () => { + // `toLlmMessages` splits the transcript on this. `createdAt` cannot do the job: it is a + // non-unique wall clock denominated in milliseconds, not in event sequence. + const { session } = makeSession() + session.append({ type: "user-message", id: "u1", text: "one" }) + session.append({ type: "turn-start", messageId: "a1" }) + session.append({ type: "text-delta", messageId: "a1", text: "two" }) + session.append({ type: "user-message", id: "u2", text: "three" }) + + const seqs = session.history().map((message) => message.startSeq) + assert.deepEqual(seqs, [1, 2, 4]) + // Monotonic, so the split is well defined. + assert.deepEqual( + [...seqs].sort((a, b) => a - b), + seqs, + ) + }) +}) diff --git a/apps/api/src/chat/ChatSession.ts b/apps/api/src/chat/ChatSession.ts index 8eda56f32..75d68426d 100644 --- a/apps/api/src/chat/ChatSession.ts +++ b/apps/api/src/chat/ChatSession.ts @@ -34,6 +34,8 @@ import { type ChatEvent, type ChatEventInput, type ChatMessage, + type ChatTaskRef, + type ChatTaskState, type ChatToolCall, type ChatTurnTenantEncoded, } from "@maple/domain/chat-session" @@ -131,6 +133,27 @@ export class ChatSession extends DurableObject> { } } + /** + * The most recent compaction, if the conversation has been summarized. + * + * A targeted reverse scan rather than a second full fold: `history()` already walks every row, + * and this is read once per turn by `toLlmMessages`. `LIMIT 1` on a descending scan stops at the + * newest compaction, which by definition is near the end of the log. + */ + compaction(): { summary: string; throughSeq: number } | undefined { + const rows = this.sql + .exec( + "SELECT seq, created_at, payload FROM events WHERE payload LIKE ? ORDER BY seq DESC LIMIT 1", + '%"type":"compaction"%', + ) + .toArray() + const row = rows[0] + if (!row) return undefined + const event = decodeChatEventPayload(row.payload, row.seq) + if (event.type !== "compaction") return undefined + return { summary: event.summary, throughSeq: event.throughSeq } + } + /** Highest assigned seq, i.e. the cursor a client that has read everything holds. */ cursor(): number { const row = this.sql.exec("SELECT MAX(seq) AS seq FROM events").one() @@ -275,7 +298,11 @@ export class ChatSession extends DurableObject> { await writer.write(encoder.encode(frameChatEvent(event))) position = event.seq } - if (events.some((event) => event.type === "turn-end")) break + // Only the *conversation's* turn ending closes the stream. A sub-agent's `turn-end` + // is tagged with `task` and merely closes its card — treating it as terminal would + // cut the connection the moment the first delegated search finished, and the rest of + // the parent's answer would only arrive on the client's next reconnect. + if (events.some((event) => event.type === "turn-end" && event.task === undefined)) break // The idle budget is spent on silence only: a batch that went out resets it, so a // long turn streams over one connection instead of being recycled mid-answer. if (!(await this.waitForAppend(SUBSCRIBE_IDLE_MS))) break @@ -405,33 +432,8 @@ export class ChatSession extends DurableObject> { * message that issued them and are completed in place by their result. */ history(): ReadonlyArray { - // A mutable mirror of `ChatMessage`. The fold concatenates deltas and completes tool calls - // in place; building it against the readonly wire type would force a copy per delta, which - // is exactly the shape that used to desynchronise `byId` from `messages`. - type Draft = { - id: string - role: ChatMessage["role"] - text: string - toolCalls: Array - createdAt: number - } - const messages: Array = [] - const byId = new Map() - - const openAssistant = (id: string, createdAt: number) => { - const existing = byId.get(id) - if (existing) return existing - const message = { - id, - role: "assistant" as const, - text: "", - toolCalls: [] as Array, - createdAt, - } - byId.set(id, message) - messages.push(message) - return message - } + const top = makeTranscript() + const nested = new Map() const rows = this.sql .exec("SELECT seq, created_at, payload FROM events ORDER BY seq ASC") @@ -439,58 +441,174 @@ export class ChatSession extends DurableObject> { for (const row of rows) { const event = decodeChatEventPayload(row.payload, row.seq) - switch (event.type) { - case "user-message": { - const message = { - id: event.id, - role: "user" as const, - text: event.text, - toolCalls: [] as Array, - createdAt: row.created_at, - } - byId.set(event.id, message) - messages.push(message) - break - } - case "turn-start": - openAssistant(event.messageId, row.created_at) - break - case "text-delta": { - // Mutate in place. Replacing the array slot with a copy left `byId` pointing at - // an object no longer in `messages`, so the *next* delta opened a brand-new - // assistant message — a 400-token reply folded into ~400 one-token messages, in - // what the browser rendered and in what the model was replayed on the next turn. - const message = openAssistant(event.messageId, row.created_at) - message.text += event.text - break - } - case "tool-call": { - const message = openAssistant(event.messageId, row.created_at) - message.toolCalls.push({ - id: event.callId, - name: event.name, - input: event.input, - ...(event.proposed === true ? { proposed: true } : {}), - } as ChatToolCall) - break - } - case "tool-result": { - const message = openAssistant(event.messageId, row.created_at) - const index = message.toolCalls.findIndex((call) => call.id === event.callId) - if (index >= 0) { - message.toolCalls[index] = { - ...message.toolCalls[index], - output: event.output, - ...(event.isError === true ? { isError: true } : {}), - } as ChatToolCall - } - break - } - case "turn-end": - break + + // A task-tagged event belongs to a sub-agent's transcript, which hangs off the parent's + // `task` tool call — not to the top-level conversation. Routing it here is what keeps a + // fan-out of sub-agents from appearing as a dozen stray assistant messages, in the + // browser *and* in what `toLlmMessages` replays to the model on the next turn. + if (event.type !== "user-message" && event.type !== "compaction" && event.task !== undefined) { + foldTaskEvent(top, nested, event, event.task, row.created_at) + continue } + + foldInto(top, event, row.created_at) } - return messages + return top.messages + } +} + +/** + * A mutable mirror of `ChatMessage`. + * + * The fold concatenates deltas and completes tool calls in place; building it against the readonly + * wire type would force a copy per delta, which is exactly the shape that used to desynchronise + * `byId` from `messages`. + */ +interface Draft { + id: string + role: ChatMessage["role"] + text: string + toolCalls: Array + createdAt: number + startSeq: number +} + +interface Transcript { + readonly messages: Array + readonly byId: Map +} + +const makeTranscript = (): Transcript => ({ messages: [], byId: new Map() }) + +const openAssistant = (transcript: Transcript, id: string, createdAt: number, startSeq: number): Draft => { + const existing = transcript.byId.get(id) + if (existing) return existing + const message: Draft = { id, role: "assistant", text: "", toolCalls: [], createdAt, startSeq } + transcript.byId.set(id, message) + transcript.messages.push(message) + return message +} + +/** + * Fold one event into a transcript. + * + * Extracted so the top-level conversation and a sub-agent's nested transcript are folded by + * *literally the same code* — two implementations of "concatenate deltas, settle tool calls in + * place" would drift, and the nested one is the harder to notice when it does. + */ +const foldInto = (transcript: Transcript, event: ChatEvent, createdAt: number): void => { + const open = (messageId: string) => openAssistant(transcript, messageId, createdAt, event.seq) + switch (event.type) { + case "user-message": { + const message: Draft = { + id: event.id, + role: "user", + text: event.text, + toolCalls: [], + createdAt, + startSeq: event.seq, + } + transcript.byId.set(event.id, message) + transcript.messages.push(message) + break + } + case "turn-start": + open(event.messageId) + break + case "text-delta": { + // Mutate in place. Replacing the array slot with a copy left `byId` pointing at an + // object no longer in `messages`, so the *next* delta opened a brand-new assistant + // message — a 400-token reply folded into ~400 one-token messages, in what the browser + // rendered and in what the model was replayed on the next turn. + open(event.messageId).text += event.text + break + } + case "tool-call": { + const message = open(event.messageId) + message.toolCalls.push({ + id: event.callId, + name: event.name, + input: event.input, + ...(event.proposed === true ? { proposed: true } : {}), + } as ChatToolCall) + break + } + case "tool-result": { + const message = open(event.messageId) + const index = message.toolCalls.findIndex((call) => call.id === event.callId) + if (index >= 0) { + message.toolCalls[index] = { + ...message.toolCalls[index], + output: event.output, + ...(event.isError === true ? { isError: true } : {}), + } as ChatToolCall + } + break + } + case "turn-retry": { + // Undo the text of the attempt that failed. The clamp is not defensive noise: + // `Stream.takeWhile(holdsTurn)` in `turn-runner.ts` can drop appends between a delta and + // the retraction that accounts for it, so `retractChars` can legitimately exceed what + // actually landed. + const message = open(event.messageId) + message.text = message.text.slice(0, Math.max(0, message.text.length - event.retractChars)) + break + } + // Inert for display. Unlike opencode — where the transcript and the model input are the same + // list — a Maple user scrolling back must still see what they actually said. Only + // `toLlmMessages` reads a compaction, through `ChatSession.compaction()`. + case "compaction": + break + case "turn-end": + break } } + +/** Terminal reason → the status the UI shows on the sub-agent's card. */ +const TASK_STATUS: Record = { + stop: "completed", + error: "error", + aborted: "aborted", + "max-steps": "aborted", +} + +/** + * Route one sub-agent event into the transcript hanging off its parent's `task` tool call. + * + * **Deny by default.** If the parent message or the parent tool call is not found, the event is + * dropped. That is the guarantee that a malformed or out-of-order child event cannot corrupt the + * parent conversation — and it costs nothing in practice, because the parent's `tool-call` + * announcement is appended strictly before the tool's `execute` runs. + */ +const foldTaskEvent = ( + top: Transcript, + /** Nested transcripts by task call id, owned by the enclosing `history()` call. */ + nested: Map, + event: ChatEvent, + ref: ChatTaskRef, + createdAt: number, +): void => { + const parent = top.byId.get(ref.parentMessageId) + if (!parent) return + const index = parent.toolCalls.findIndex((call) => call.id === ref.id) + if (index < 0) return + + const child = nested.get(ref.id) ?? makeTranscript() + nested.set(ref.id, child) + foldInto(child, event, createdAt) + + // `ChatToolCall` is the readonly wire type, so the call is rebuilt rather than mutated. Keying + // the nested transcript by task id rather than by object identity is what makes that safe. + parent.toolCalls[index] = { + ...parent.toolCalls[index]!, + task: { + id: ref.id, + agent: ref.agent, + status: event.type === "turn-end" ? (TASK_STATUS[event.reason] ?? "completed") : "running", + // The nested drafts are structurally `ChatSubMessage` already — a sub-agent cannot nest + // further, so no `task` field is ever present on them. + messages: child.messages, + }, + } as ChatToolCall +} diff --git a/apps/api/src/chat/agent.test.ts b/apps/api/src/chat/agent.test.ts index 48e002126..0437b30c1 100644 --- a/apps/api/src/chat/agent.test.ts +++ b/apps/api/src/chat/agent.test.ts @@ -7,9 +7,13 @@ */ import { assert, describe, it } from "vitest" import { Effect, Layer, Stream } from "effect" -import { LLM, LLMClient, LLMEvent, type LLMRequest, type Model } from "@maple/llm" +import { LLMClient, LLMEvent, type LLMRequest, type Model } from "@maple/llm" import { CloudflareWorkersAI } from "@maple/llm/providers/cloudflare" import { runChatTurn, type ChatTurnEvent } from "./agent" +import { MAX_STEP_ATTEMPTS } from "./llm-retry" +import { DEFAULT_RULESET } from "./permissions" +import type { AgentDefinition } from "./agents" +import { PermissionRule } from "@maple/domain/permission" import type { TenantContext } from "@/services/auth/tenant-context" const TENANT: TenantContext = { @@ -38,53 +42,88 @@ const toolCall = (id: string, name: string): LLMEvent => * `stream` is the only method the turn uses; `prepare`/`generate` are present because the service * interface has them and a partial stub would be a lie about what is being exercised. */ +type Failure = { + readonly events: ReadonlyArray + readonly fail: true + /** + * Vendored reason tag. Defaults to a retryable `ProviderInternal`; `"Authentication"` (or any + * other terminal reason) is how a test asserts the non-retrying path. Note that `Transport` and + * `InvalidProviderOutput` carry `retryable: false` on the wire and are still retried — see + * `llm-retry.ts`. + */ + readonly reason?: string + readonly retryable?: boolean +} + type Step = /** A clean step: these events, then the stream completes. */ | ReadonlyArray /** The stream fails after emitting `events` — the realistic partial-stream case. */ - | { readonly events: ReadonlyArray; readonly fail: true } + | Failure + +/** Every request the turn issued, in order. Lets a test assert on tools, toolChoice and messages. */ +type RequestLog = Array -const stubModel = (steps: ReadonlyArray) => { +const stubModel = (steps: ReadonlyArray, log: RequestLog = []) => { let step = 0 const service = { prepare: () => Effect.die(new Error("prepare is not used by runChatTurn")), generate: () => Effect.die(new Error("generate is not used by runChatTurn")), - stream: (_request: LLMRequest) => { + stream: (request: LLMRequest) => { + log.push(request) const scripted = steps[step] ?? [finish()] step += 1 + if (Array.isArray(scripted)) return Stream.fromIterable(scripted as ReadonlyArray) + const partial = scripted as Failure // The vendored error shape the turn maps through `toLlmCallError`. const failure = { _tag: "LLMError", module: "test", method: "stream", - reason: { _tag: "ProviderInternal" }, + reason: { _tag: partial.reason ?? "ProviderInternal" }, message: "upstream exploded", - retryable: true, + retryable: partial.retryable ?? partial.reason === undefined, } as never - if (Array.isArray(scripted)) return Stream.fromIterable(scripted as ReadonlyArray) - const partial = scripted as { events: ReadonlyArray } return Stream.concat(Stream.fromIterable(partial.events), Stream.fail(failure)) }, } - return Layer.succeed(LLMClient.Service, service as never) + return { layer: Layer.succeed(LLMClient.Service, service as never), log, calls: () => step } } const collect = ( steps: ReadonlyArray, - overrides: { readonly sessionId?: string; readonly isCurrent?: () => boolean } = {}, -) => - runChatTurn({ + overrides: { + readonly sessionId?: string + readonly isCurrent?: () => boolean + readonly agent?: AgentDefinition + } = {}, +) => { + const stub = stubModel(steps) + return runChatTurn({ sessionId: overrides.sessionId ?? "org_test:tab", tenant: TENANT, model: MODEL, messages: [], messageId: "m1", ...(overrides.isCurrent ? { isCurrent: overrides.isCurrent } : {}), + ...(overrides.agent ? { agent: overrides.agent } : {}), }).pipe( Stream.runCollect, Effect.map((events) => Array.from(events) as ChatTurnEvent[]), - Effect.provide(stubModel(steps)), + Effect.provide(stub.layer), + Effect.map((events) => ({ events, requests: stub.log, calls: stub.calls() })), ) +} + +/** The common case: only the events matter. */ +const collectEvents = ( + steps: ReadonlyArray, + overrides: { + readonly sessionId?: string + readonly isCurrent?: () => boolean + readonly agent?: AgentDefinition + } = {}, +) => collect(steps, overrides).pipe(Effect.map((result) => result.events)) const types = (events: ReadonlyArray) => events.map((event) => event.type) @@ -92,7 +131,7 @@ const terminal = (events: ReadonlyArray) => events.filter((event) describe("runChatTurn", () => { it("emits turn-start, the text, and exactly one turn-end", async () => { - const events = await Effect.runPromise(collect([[textDelta("Hello"), finish()]])) + const events = await Effect.runPromise(collectEvents([[textDelta("Hello"), finish()]])) assert.deepEqual(types(events), ["turn-start", "text-delta", "turn-end"]) assert.lengthOf(terminal(events), 1) @@ -100,7 +139,7 @@ describe("runChatTurn", () => { it("coalesces adjacent text deltas into one event without losing any text", async () => { const chunks = ["Check", "ing ", "the ", "traces", "."] - const events = await Effect.runPromise(collect([[...chunks.map(textDelta), finish()]])) + const events = await Effect.runPromise(collectEvents([[...chunks.map(textDelta), finish()]])) // One durable row, one SSE frame and one React commit per token is more fidelity than a // screen can show, and the transcript render cost is paid per commit. @@ -114,7 +153,7 @@ describe("runChatTurn", () => { it("keeps text ahead of the tool calls it precedes", async () => { const events = await Effect.runPromise( - collect([ + collectEvents([ [textDelta("Looking"), textDelta(" it up"), toolCall("c1", "create_alert_rule"), finish()], ]), ) @@ -126,8 +165,10 @@ describe("runChatTurn", () => { assert.equal(delta?.type === "text-delta" ? delta.text : undefined, "Looking it up") }) - it("emits exactly ONE turn-end when the model stream fails", async () => { - const events = await Effect.runPromise(collect([{ events: [textDelta("part")], fail: true }])) + it("emits exactly ONE turn-end when the model stream fails terminally", async () => { + const events = await Effect.runPromise( + collectEvents([{ events: [textDelta("part")], fail: true, reason: "Authentication" }]), + ) // The regression: `Stream.concat`'s second half ran unconditionally, so a failed stream that // still assembled a partial response emitted a second terminal event after the error one. @@ -142,7 +183,13 @@ describe("runChatTurn", () => { // A partial stream that carries a tool call and then dies: `LLMResponse.fromEvents` will // happily assemble it, which is exactly how the second half used to run past the error. const events = await Effect.runPromise( - collect([{ events: [toolCall("c1", "find_errors"), textDelta("partial")], fail: true }]), + collectEvents([ + { + events: [toolCall("c1", "find_errors"), textDelta("partial")], + fail: true, + reason: "Authentication", + }, + ]), ) assert.isEmpty( @@ -152,7 +199,9 @@ describe("runChatTurn", () => { }) it("stops on an approval-gated tool with a proposal and no result", async () => { - const events = await Effect.runPromise(collect([[toolCall("c1", "create_alert_rule"), finish()]])) + const events = await Effect.runPromise( + collectEvents([[toolCall("c1", "create_alert_rule"), finish()]]), + ) const proposals = events.filter((event) => event.type === "tool-call") assert.lengthOf(proposals, 1) @@ -168,7 +217,7 @@ describe("runChatTurn", () => { it("stops without a second terminal event when the turn is superseded", async () => { const events = await Effect.runPromise( - collect([[textDelta("hi"), finish()]], { isCurrent: () => false }), + collectEvents([[textDelta("hi"), finish()]], { isCurrent: () => false }), ) // An abort already recorded the terminal event on the session; emitting another here would @@ -176,3 +225,246 @@ describe("runChatTurn", () => { assert.isEmpty(terminal(events)) }) }) + +describe("runChatTurn max steps", () => { + it("spends a final tool-less step answering instead of stopping dead", async () => { + // Ten steps that each call a read-only tool, then a text-only step. + const looping: Step = [toolCall("c1", "find_errors"), finish()] + const result = await Effect.runPromise( + collect([ + ...Array.from({ length: 10 }, () => looping), + [textDelta("Here is what I found."), finish()], + ]), + ) + + // The turn used to end here on a wall of tool rows with no words at all. + const last = result.events[result.events.length - 2] + assert.equal(last?.type, "text-delta") + assert.equal(last?.type === "text-delta" ? last.text : undefined, "Here is what I found.") + + // The closing step cannot loop even if the model ignores the instruction. + const closing = result.requests[result.requests.length - 1] + assert.isEmpty(closing?.tools ?? [{}]) + assert.equal(closing?.toolChoice?.type, "none") + + // The signal the client badges on survives — it just arrives with an answer attached now. + assert.lengthOf(terminal(result.events), 1) + const end = terminal(result.events)[0] + assert.equal(end?.type === "turn-end" ? end.reason : undefined, "max-steps") + }, 30_000) +}) + +const retries = (events: ReadonlyArray) => + events.filter((event) => event.type === "turn-retry") + +const textOf = (events: ReadonlyArray) => + events + .filter((event) => event.type === "text-delta") + .map((event) => (event.type === "text-delta" ? event.text : "")) + .join("") + +/** What a reader ends up with — the same concatenate-and-retract fold `ChatSession.history()` does. */ +const fold = (events: ReadonlyArray): string => + events.reduce((text, event) => { + if (event.type === "text-delta") return text + event.text + if (event.type === "turn-retry") return text.slice(0, Math.max(0, text.length - event.retractChars)) + return text + }, "") + +describe("runChatTurn retry", () => { + it("retracts nothing when the failed attempt's batch never flushed", async () => { + const result = await Effect.runPromise( + collect([{ events: [textDelta("Hello wo")], fail: true }, [textDelta("Hello world."), finish()]]), + ) + + // `Stream.groupedWithin` drops its pending buffer on an upstream failure rather than + // flushing it, so a provider that dies inside one batching window emitted nothing and there + // is nothing to take back. The marker is still emitted — it is also the progress signal. + const retracted = retries(result.events) + assert.lengthOf(retracted, 1) + assert.equal(retracted[0]?.type === "turn-retry" ? retracted[0].retractChars : undefined, 0) + assert.equal(result.calls, 2) + assert.equal(textOf(result.events), "Hello world.") + assert.lengthOf(terminal(result.events), 1) + }) + + it("retracts exactly the text the failed attempt did flush", async () => { + // A full batch (`DELTA_BATCH_SIZE`) flushes on the size cap regardless of timing, so this is + // the deterministic stand-in for a real provider stream, where the 16ms window fires + // constantly and most text has already shipped by the time the body dies. + const flushed = Array.from({ length: 24 }, (_, i) => textDelta(String(i % 10))) + const result = await Effect.runPromise( + collect([ + { events: [...flushed, textDelta("buffered")], fail: true }, + [textDelta("the real answer"), finish()], + ]), + ) + + // The test that catches duplicated text: without the retraction the durable fold reads the + // abandoned prefix followed by the retry's text — permanently, because deltas concatenate + // and the log is durable. + const retracted = retries(result.events) + assert.lengthOf(retracted, 1) + const marker = retracted[0] + assert.equal(marker?.type === "turn-retry" ? marker.retractChars : undefined, 24) + assert.equal(marker?.type === "turn-retry" ? marker.attempt : undefined, 2) + + // Fold the whole event list the way `ChatSession.history()` does, and check the reader lands + // on the retry's text alone — the abandoned prefix is gone, not doubled. + assert.equal(fold(result.events), "the real answer") + }) + + it("retries a Transport failure even though it reports retryable: false", async () => { + // The whole point of `llm-retry.ts`. `TransportReason` and `InvalidProviderOutputReason` + // hardcode `retryable = false`, and they are exactly how a body that dies mid-stream + // surfaces — a classifier that trusted the flag would pass a naive test and do nothing. + for (const reason of ["Transport", "InvalidProviderOutput"]) { + const result = await Effect.runPromise( + collect([{ events: [], fail: true, reason, retryable: false }, [textDelta("ok"), finish()]]), + ) + assert.equal(result.calls, 2, `${reason} should have been retried`) + assert.equal(textOf(result.events), "ok") + } + }) + + it("does not retry a failure that would fail identically", async () => { + const result = await Effect.runPromise( + collect([{ events: [], fail: true, reason: "Authentication" }, [textDelta("never"), finish()]]), + ) + + assert.equal(result.calls, 1) + assert.isEmpty(retries(result.events)) + const end = terminal(result.events)[0] + assert.equal(end?.type === "turn-end" ? end.reason : undefined, "error") + }) + + it("gives up after the attempt budget and ends the turn once", async () => { + const failing: Step = { events: [], fail: true } + const result = await Effect.runPromise( + collect([failing, failing, failing, failing, failing, failing]), + ) + + assert.equal(result.calls, MAX_STEP_ATTEMPTS) + assert.lengthOf(retries(result.events), MAX_STEP_ATTEMPTS - 1) + assert.lengthOf(terminal(result.events), 1) + const end = terminal(result.events)[0] + assert.equal(end?.type === "turn-end" ? end.reason : undefined, "error") + }, 30_000) + + it("stops during backoff without emitting a terminal event", async () => { + // `isCurrent` is re-checked after the sleep, so an abort landing mid-backoff wins. The DO + // already wrote the terminal event when it cleared the claim. + let current = true + const result = await Effect.runPromise( + collect([{ events: [textDelta("partial")], fail: true }, [textDelta("never"), finish()]], { + isCurrent: () => { + const value = current + // Still current when the stream fails (so the retry is scheduled), superseded by + // the time the backoff elapses. + current = false + return value + }, + }), + ) + + assert.equal(result.calls, 1) + assert.isEmpty(terminal(result.events)) + }) +}) + +const agentWith = (overrides: Partial): AgentDefinition => ({ + name: "test", + description: "test", + mode: "primary", + prompt: "be brief", + permission: DEFAULT_RULESET, + ...overrides, +}) + +const toolNames = (request: LLMRequest | undefined) => (request?.tools ?? []).map((tool) => tool.name) + +describe("runChatTurn permissions", () => { + it("never offers the model a denied tool", () => { + // Stronger and cheaper than refusing the call afterwards: a tool the model cannot see is a + // tool it cannot be talked into trying. + const ruleset = [ + new PermissionRule({ tool: "*", action: "allow" }), + new PermissionRule({ tool: "create_*", action: "deny" }), + ] + return Effect.runPromise( + collect([[textDelta("ok"), finish()]], { agent: agentWith({ permission: ruleset }) }).pipe( + Effect.map((result) => { + const offered = toolNames(result.requests[0]) + assert.notInclude(offered, "create_alert_rule") + assert.include(offered, "find_errors") + }), + ), + ) + }) + + it("still offers an approval-gated tool, and still stops on it", async () => { + // `ask` is visible-but-interrupting. Hiding it would make the model unable to propose the + // mutation at all, which is not what approval means. + const result = await Effect.runPromise(collect([[toolCall("c1", "create_alert_rule"), finish()]])) + + assert.include(toolNames(result.requests[0]), "create_alert_rule") + const proposals = result.events.filter((event) => event.type === "tool-call") + assert.lengthOf(proposals, 1) + assert.equal(proposals[0]?.type === "tool-call" ? proposals[0].proposed : undefined, true) + }) + + it("executes a mutation when the ruleset allows it outright", async () => { + // The capability rulesets unlock: an agent that is trusted with a tool no longer has to + // round-trip through an approval card for it. + const ruleset = [new PermissionRule({ tool: "*", action: "allow" })] + const result = await Effect.runPromise( + collect( + [ + [toolCall("c1", "create_alert_rule"), finish()], + [textDelta("done"), finish()], + ], + { + agent: agentWith({ permission: ruleset }), + }, + ), + ) + + const announced = result.events.filter((event) => event.type === "tool-call") + assert.equal(announced[0]?.type === "tool-call" ? announced[0].proposed : undefined, undefined) + assert.isNotEmpty(result.events.filter((event) => event.type === "tool-result")) + }) + + it("honours an agent's own step budget", async () => { + const looping: Step = [toolCall("c1", "find_errors"), finish()] + const result = await Effect.runPromise( + collect([...Array.from({ length: 6 }, () => looping), [textDelta("summary"), finish()]], { + agent: agentWith({ steps: 3 }), + }), + ) + + // Three tool-calling steps, then the tool-less closing one. + assert.equal(result.calls, 4) + const end = terminal(result.events)[0] + assert.equal(end?.type === "turn-end" ? end.reason : undefined, "max-steps") + }, 30_000) +}) + +describe("runChatTurn max steps, defensively", () => { + it("ends rather than looping when a provider calls tools on the closing step", async () => { + // The closing step is sent with `tools: []` and `toolChoice: "none"`. A provider that emits a + // call anyway must not get another closing step, or `MAX_STEPS` stops being a bound at all. + const looping: Step = [toolCall("c1", "find_errors"), finish()] + const result = await Effect.runPromise( + collect( + Array.from({ length: 20 }, () => looping), + { agent: agentWith({ steps: 2 }) }, + ), + ) + + // Two tool-calling steps, one closing step, then stop — not twenty. + assert.equal(result.calls, 3) + assert.lengthOf(terminal(result.events), 1) + const end = terminal(result.events)[0] + assert.equal(end?.type === "turn-end" ? end.reason : undefined, "max-steps") + }, 30_000) +}) diff --git a/apps/api/src/chat/agent.ts b/apps/api/src/chat/agent.ts index 08ed93144..08f3c243f 100644 --- a/apps/api/src/chat/agent.ts +++ b/apps/api/src/chat/agent.ts @@ -16,10 +16,11 @@ * user's action, authenticated as the user, which is exactly where it belongs. */ import { - chatModeFromSessionId, investigationIdFromChatSessionId, type ChatEvent, + type ChatTaskRef, } from "@maple/domain/chat-session" +import { evaluatePermission, type PermissionRuleset } from "@maple/domain/permission" import { AiTriageResult, SubmitDiagnosisRequest } from "@maple/domain/http" import { InvestigationId } from "@maple/domain/primitives" import { @@ -33,20 +34,44 @@ import { type Usage, } from "@maple/llm" import { Tool, ToolFailure, ToolRuntime, toDefinitions, type Tools } from "@maple/llm" -import { Cause, Effect, Option, Schema, Stream } from "effect" -import { toLlmCallError } from "@/platform/Llm" +import { Duration, Effect, Option, Schema, Stream } from "effect" +import { contextLimitOf, outputLimitOf, toLlmCallError } from "@/platform/Llm" import type { TenantContext } from "@/services/auth/tenant-context" -import { callMcpTool } from "@/mcp/dispatcher" -import { CurrentMcpTenant } from "@/mcp/lib/query-warehouse" -import { MUTATING_TOOL_NAMES } from "@/mcp/tools/mutating" -import { mapleToolDefinitions, toInputSchema } from "@/mcp/tools/registry" -import { buildSystemPrompt } from "./modes" + +import { buildMapleTools, summarizeToolFailure, withRuntimeServices } from "@/mcp/tools/llm-tools" +import { isNearContextLimit, pruneToolResults } from "./context-budget" +import { agentForSession, buildSystemPrompt, spawnableFor, type AgentDefinition } from "./agents" +import { buildTaskTool, hasStepBudget, makeTaskBudget, type TaskBudget } from "./task-tool" +import { + isRetryableStepFailure, + makeStepRetryBudget, + stepRetryDelayMs, + MAX_STEP_ATTEMPTS, + STEP_RETRY_BUDGET_MS, + type StepRetryBudget, +} from "./llm-retry" const decodeInvestigationIdOption = Schema.decodeUnknownOption(InvestigationId) -/** Hard cap on assistant turns per submission. */ +/** + * Hard cap on *tool-calling* assistant turns per submission. + * + * A turn that hits it gets one further, tool-less step so the model can answer from what it found; + * see `MAX_STEPS_NOTICE`. + */ const MAX_STEPS = 10 +/** + * What the model is told when it runs out of steps. + * + * The turn used to stop dead here, so the user was left with a wall of tool rows and no words — + * the model had gathered the answer and never got to say it. + */ +const MAX_STEPS_NOTICE = + "You have reached the maximum number of tool calls for this turn. Do NOT call any more tools. " + + "Using only what you have already found, give the user your answer now, and say plainly what " + + "you could not determine." + /** Fan-out cap for tool calls issued in the same assistant turn. */ const TOOL_CONCURRENCY = 4 @@ -87,44 +112,6 @@ const addUsage = (total: TurnUsage, usage: Usage | undefined): void => { total.cacheRead += usage?.cacheReadInputTokens ?? 0 } -/** - * Re-pin the service requirements of an MCP tool handler. See the identical helper in - * `workflows/triage-agent.ts` — `MapleToolDefinition.handler` erases its requirements to `any` - * at the registry boundary, and `Tool.make` insists on `never`. - */ -// eslint-disable-next-line @typescript-eslint/no-explicit-any -const withRuntimeServices = (effect: Effect.Effect): Effect.Effect => - effect as Effect.Effect - -const toolResultText = (result: { content: ReadonlyArray<{ text: string }> }): string => - result.content.map((block) => block.text).join("\n") - -/** - * A one-line reason for a failed tool, safe to hand the model. - * - * `String(cause)` renders the whole Effect cause: stack frames, and — inside a `DatabaseError` — - * connection details. That went into the model's context and, through the tool-result event, into - * a durable transcript the browser reads back. - */ -const summarizeToolFailure = (cause: Cause.Cause): string => { - const failure = cause.reasons.find(Cause.isFailReason) - const error: unknown = failure?.error - if (error instanceof Error) return error.message - if (error && typeof error === "object" && "message" in error) { - const message = (error as { message?: unknown }).message - if (typeof message === "string") return message - } - return "the tool failed" -} - -/** - * Description suffix on gated tools. The model still calls them normally; it just needs to know - * the call is a proposal so it stops rather than narrating a completed change. - */ -const APPROVAL_NOTE = - "\n\nThis is an approval-gated action. Calling it proposes the change for the user to approve; " + - "it does NOT take effect until they do. Call it once with the intended arguments and stop." - export interface ChatTurnInput { readonly sessionId: string readonly tenant: TenantContext @@ -148,6 +135,32 @@ export interface ChatTurnInput { readonly isCurrent?: () => boolean /** Accumulates this turn's token usage; see {@link TurnUsage}. */ readonly usage?: TurnUsage + /** + * Which agent this turn runs as. Defaults to the primary agent the session id names, so + * existing callers keep the behaviour their mode already had. + */ + readonly agent?: AgentDefinition + /** + * Set when this turn is a sub-agent nested inside a parent turn: every event it produces is + * stamped with this ref, which is what routes them into the parent's task card rather than the + * top-level conversation. + */ + readonly task?: ChatTaskRef + /** Nesting depth. 0 is the conversation's own turn. */ + readonly depth?: number + /** Shared across the parent and every descendant; see {@link TaskBudget}. */ + readonly taskBudget?: TaskBudget + /** + * Sink for events produced by a nested sub-agent turn. + * + * A side channel rather than a merge into the returned stream. `Stream.merge` would race the + * parent's terminal `turn-end` against undrained child events, and `ChatSession.pump` closes + * the SSE connection on a terminal event — so a child's tail could be written to the durable log + * *after* the turn had been declared over. The consumer's `Stream.runForEach` is strictly + * sequential, so an event emitted from inside a tool's `execute` is deterministically ordered + * after the tool-call announcement that introduced it and before the tool-result that closes it. + */ + readonly emit?: (event: ChatTurnEvent) => void } /** @@ -223,51 +236,20 @@ export const buildSubmitDiagnosisTool = ( return tools } -/** All Maple tools, with mutating ones flagged. Read-only tools execute; gated ones never do. */ -const buildChatTools = (tenant: TenantContext): Tools => - Object.fromEntries( - mapleToolDefinitions.map((definition) => { - const gated = MUTATING_TOOL_NAMES.has(definition.name) - return [ - definition.name, - Tool.make({ - description: gated ? `${definition.description}${APPROVAL_NOTE}` : definition.description, - jsonSchema: toInputSchema(definition.schema), - // A gated tool still executes as far as `Tool.make` is concerned, but the loop - // never dispatches it — it breaks on the proposal first. Keeping a real handler - // here (rather than omitting it) means the tool schema the model sees is - // identical to the read-only case, and `POST /api/chat/apply` remains the only - // path that actually mutates. - execute: (params): Effect.Effect => - gated - ? Effect.fail( - new ToolFailure({ - message: `${definition.name} requires user approval and was not executed.`, - }), - ) - : withRuntimeServices( - callMcpTool(definition.name, params).pipe( - Effect.provideService(CurrentMcpTenant, tenant), - Effect.flatMap((result) => - result.isError - ? Effect.fail( - new ToolFailure({ message: toolResultText(result) }), - ) - : Effect.succeed(toolResultText(result)), - ), - Effect.catchCause((cause) => - Effect.fail( - new ToolFailure({ - message: `Tool failed: ${summarizeToolFailure(cause)}`, - }), - ), - ), - ), - ), - }), - ] - }), - ) +/** + * All Maple tools, with mutating ones gated. + * + * A gated tool still carries a real handler (rather than being omitted) so the schema the model + * sees is identical to the read-only case — but the loop never dispatches it, because it breaks on + * the proposal first, and `POST /api/chat/apply` remains the only path that actually mutates. + */ +const buildChatTools = (tenant: TenantContext, ruleset: PermissionRuleset): Tools => + buildMapleTools(tenant, { + // `deny` means the model never sees the tool. That is a stronger guarantee than refusing the + // call afterwards, and it is free — an unoffered tool cannot be called. + include: (name) => evaluatePermission(ruleset, name) !== "deny", + gate: (name) => evaluatePermission(ruleset, name) === "ask", + }) /** Distributive `Omit`, so each union member keeps its own shape. */ type WithoutSeq = T extends unknown ? Omit : never @@ -290,30 +272,77 @@ export type ChatTurnEvent = WithoutSeq => Stream.unwrap( Effect.sync(() => { - const tools = { ...buildChatTools(input.tenant), ...input.extraTools } - const system = buildSystemPrompt({ mode: chatModeFromSessionId(input.sessionId) }) + const agent = input.agent ?? agentForSession(input.sessionId) + const taskBudget = input.taskBudget ?? makeTaskBudget() + const tools = { + ...buildChatTools(input.tenant, agent.permission), + // Delegation is opt-in per agent: an agent with no `spawns` never sees `task` at all. + ...buildTaskTool(input, spawnableFor(agent), taskBudget, runChatTurn), + ...input.extraTools, + } const request = LLM.request({ id: input.messageId, model: input.model, - system, + system: buildSystemPrompt(agent), messages: [...input.messages], tools: toDefinitions(tools), }) - const start: ChatTurnEvent = { type: "turn-start", messageId: input.messageId } - return Stream.concat(Stream.fromIterable([start]), runStep(input, tools, request, 0)) + const start = tagged(input, { type: "turn-start", messageId: input.messageId }) + const state: StepState = { + step: 0, + attempt: 0, + budget: makeStepRetryBudget(), + agent, + taskBudget, + } + return Stream.concat(Stream.fromIterable([start]), runStep(input, tools, request, state)) }), ) +/** + * Where one step sits in the turn. + * + * `attempt` is 0-based and resets per step; `budget` is shared across the whole turn, because a + * per-step cap composes badly over `MAX_STEPS` steps — ten steps each retrying three times is + * minutes of pure backoff, and the DO's turn slot is held the entire time. + */ +interface StepState { + readonly step: number + readonly attempt: number + readonly budget: StepRetryBudget + /** Resolved once at the top of the turn, so every step agrees on the ruleset and step cap. */ + readonly agent: AgentDefinition + /** Shared with every descendant sub-agent turn. */ + readonly taskBudget: TaskBudget + /** + * This is the tool-less step that closes a turn which ran out of steps. Its natural exit is + * "no tool calls", which would otherwise report `"stop"` and lose the signal the client badges + * on — so the reason is carried here instead. + */ + readonly closing?: true +} + +/** + * Stamp an event with the ref that routes it into a parent's task card. + * + * Every emission site goes through this, so a sub-agent's events can never leak into the top-level + * conversation by omission — the tag is applied once, at the boundary, rather than remembered at + * each of the seven places a turn emits. + */ +const tagged = (input: ChatTurnInput, event: E): E => + input.task === undefined ? event : { ...event, task: input.task } + const turnEnd = ( input: ChatTurnInput, reason: Extract["reason"], error?: string, -): ChatTurnEvent => ({ - type: "turn-end", - messageId: input.messageId, - reason, - ...(error === undefined ? {} : { error }), -}) +): ChatTurnEvent => + tagged(input, { + type: "turn-end", + messageId: input.messageId, + reason, + ...(error === undefined ? {} : { error }), + }) /** A turn with no session attached (tests, one-shot callers) is always current. */ const isCurrent = (input: ChatTurnInput): boolean => input.isCurrent === undefined || input.isCurrent() @@ -328,9 +357,14 @@ const runStep = ( input: ChatTurnInput, tools: Tools, request: LLMRequest, - step: number, + state: StepState, ): Stream.Stream => Stream.suspend(() => { + // Counted per model call, not per logical step, because a retry costs the same wall clock + // and the same money. Shared with every descendant, so a fan-out of sub-agents cannot + // multiply its way past the turn's ceiling. + state.taskBudget.stepsUsed += 1 + const collected: Array = [] // Set by the catch below. `Stream.concat`'s second half runs unconditionally, so without an // explicit flag a failed stream that still assembled a partial response would emit a @@ -339,6 +373,15 @@ const runStep = ( // Those extra events land invisibly (the SSE route stops at the first `turn-end`) and // surface on the next reload. let failed = false + // Characters of *this attempt's* text that reached the log. Counted after the batching + // window, not at the raw delta, so it matches exactly what the consumer appended — that + // equality is what makes `retractChars` a complete undo rather than an approximation. + // + // It also means a batch still buffering when the stream fails contributes nothing, because + // `Stream.groupedWithin` discards its pending buffer on an upstream failure rather than + // flushing it. So a provider that dies inside one batching window costs a retraction of + // zero, and only text that actually reached a consumer is ever taken back. + let emitted = 0 const live: Stream.Stream = LLM.stream(request).pipe( Stream.tap((event) => Effect.sync(() => collected.push(event))), @@ -350,21 +393,87 @@ const runStep = ( // tool calls and the terminal event live in the concatenated segment below, so nothing // here can reorder them. Stream.groupedWithin(DELTA_BATCH_SIZE, DELTA_BATCH_WINDOW), - Stream.map( - (events): ChatTurnEvent => ({ - type: "text-delta", - messageId: input.messageId, - text: events.map((event) => ("text" in event ? event.text : "")).join(""), - }), - ), - // A model failure ends the turn as a recorded event rather than killing the stream — - // the session log is durable, so a client reconnecting after the failure must still be - // able to read why the turn stopped. + Stream.map((events): ChatTurnEvent => { + const text = events.map((event) => ("text" in event ? event.text : "")).join("") + emitted += text.length + return tagged(input, { type: "text-delta", messageId: input.messageId, text }) + }), + // A model failure either retries the step or ends the turn as a recorded event. Either + // way it does not kill the stream: the session log is durable, so a client reconnecting + // after the failure must still be able to read what happened. Stream.catch((error) => { failed = true - return Stream.fromIterable([ - turnEnd(input, "error", toLlmCallError("chat.turn", error).message), - ]) + const called = toLlmCallError("chat.turn", error) + + // Aborted mid-stream. The DO already recorded the terminal event when it cleared the + // claim, so emitting anything here would be a second one. + if (!isCurrent(input)) return Stream.empty + + // Overflow is the one failure worth retrying with a *different* request. Sending the + // same oversized transcript again cannot start fitting, so `isRetryableStepFailure` + // refuses it — but a pruned transcript is a genuinely new attempt. This is what + // `LlmCallError.contextOverflow` was added for; nothing acted on it before. + if (called.contextOverflow) { + const pruned = pruneToolResults(request) + if (pruned === request || state.attempt + 1 >= MAX_STEP_ATTEMPTS) { + return Stream.fromIterable([turnEnd(input, "error", called.message)]) + } + return Stream.concat( + Stream.fromIterable([ + tagged(input, { + type: "turn-retry" as const, + messageId: input.messageId, + attempt: state.attempt + 2, + retractChars: emitted, + reason: called.reason, + delayMs: 0, + }), + ]), + runStep(input, tools, pruned, { ...state, attempt: state.attempt + 1 }), + ) + } + + const delayMs = stepRetryDelayMs(state.attempt) + const affordable = state.budget.spentMs + delayMs <= STEP_RETRY_BUDGET_MS + if ( + !isRetryableStepFailure(called) || + state.attempt + 1 >= MAX_STEP_ATTEMPTS || + !affordable + ) { + return Stream.fromIterable([turnEnd(input, "error", called.message)]) + } + state.budget.spentMs += delayMs + + // The retraction and the progress signal are one event: either alone is useless. + // Safe to express as a character count because a failed attempt emitted nothing but + // text — tool calls live in `settleAndRecurse`, which `failed` short-circuits. + const marker = tagged(input, { + type: "turn-retry" as const, + messageId: input.messageId, + attempt: state.attempt + 2, + retractChars: emitted, + reason: called.reason, + delayMs, + }) + return Stream.concat( + Stream.fromIterable([marker]), + // `Stream.unwrap` + `Effect.sleep` rather than `Stream.retry`: a schedule would + // resubscribe this whole pipeline, replaying the deltas it already emitted, and + // would leave nowhere to put the retraction between attempts. + Stream.unwrap( + Effect.sleep(Duration.millis(delayMs)).pipe( + Effect.map(() => + // Re-checked *after* the sleep: an abort landing during backoff wins. + isCurrent(input) + ? runStep(input, tools, request, { + ...state, + attempt: state.attempt + 1, + }) + : Stream.empty, + ), + ), + ), + ) }), ) @@ -386,28 +495,38 @@ const runStep = ( .filter(LLMEvent.is.toolCall) .filter((call) => !call.providerExecuted) - if (calls.length === 0) return Stream.fromIterable([turnEnd(input, "stop")]) + if (calls.length === 0) { + return Stream.fromIterable([turnEnd(input, state.closing ? "max-steps" : "stop")]) + } + + // The closing step is sent with `tools: []` and `toolChoice: "none"`, so a call here + // means the provider ignored both. Ending rather than dispatching keeps `MAX_STEPS` + // a real bound: without this the closing step would recurse into another closing + // step, and a provider that always emits a call would loop forever. + if (state.closing) return Stream.fromIterable([turnEnd(input, "max-steps")]) // The real interrupt. A gated call ends the turn immediately — the client renders an // approval card from this event and applies it through `POST /api/chat/apply`. // Read-only calls issued in the same turn are dropped rather than half-run, so the // transcript never shows a partial turn. - const gated = calls.find((call) => MUTATING_TOOL_NAMES.has(call.name)) + const gated = calls.find( + (call) => evaluatePermission(state.agent.permission, call.name) === "ask", + ) if (gated) { - const proposal: ChatTurnEvent = { - type: "tool-call", + const proposal = tagged(input, { + type: "tool-call" as const, messageId: input.messageId, callId: gated.id, name: gated.name, input: gated.input, proposed: true, - } + }) return Stream.fromIterable([proposal, turnEnd(input, "stop")]) } - const announced = calls.map( - (call): ChatTurnEvent => ({ - type: "tool-call", + const announced = calls.map((call) => + tagged(input, { + type: "tool-call" as const, messageId: input.messageId, callId: call.id, name: call.name, @@ -430,9 +549,9 @@ const runStep = ( { concurrency: TOOL_CONCURRENCY }, ) - const results = dispatched.map( - ([call, outcome]): ChatTurnEvent => ({ - type: "tool-result", + const results = dispatched.map(([call, outcome]) => + tagged(input, { + type: "tool-result" as const, messageId: input.messageId, callId: call.id, output: outcome.result.value, @@ -444,28 +563,75 @@ const runStep = ( // transcript is not left with dangling calls, then stop. if (!isCurrent(input)) return Stream.fromIterable(results) - if (step + 1 >= MAX_STEPS) { - return Stream.fromIterable([...results, turnEnd(input, "max-steps")]) + const transcript = [ + ...request.messages, + response.message, + ...dispatched.map(([call, outcome]) => + Message.tool( + ToolResultPart.make({ + id: call.id, + name: call.name, + result: outcome.result, + }), + ), + ), + ] + + /** + * Prune before the next step if the *provider's own* count says we are near + * the wall. Acting on the reported figure rather than an estimate is what + * makes this trustworthy — the estimate exists only to decide whether a + * prune is worth doing. + */ + const withBudget = (next: LLMRequest): LLMRequest => + isNearContextLimit(response.usage?.inputTokens ?? 0, { + context: contextLimitOf(input.model), + output: outputLimitOf(input.model), + }) + ? pruneToolResults(next) + : next + + // Out of steps. Rather than cutting the turn off after a wall of tool rows + // with no words — which is what the user was left with — spend one more + // non-tool step letting the model answer from what it already found. + // + // A trailing *user* instruction, not opencode's assistant prefill: prefill is + // an Anthropic-shaped affordance, and Maple's default route is OpenRouter. + // `tools: []` and `toolChoice: "none"` together mean the closing step cannot + // loop even if the model ignores the instruction, so `MAX_STEPS` keeps + // meaning "at most this many tool-calling steps". + // Either this turn's own step cap, or the budget shared with every sub-agent it + // spawned. Both land in the same place: one closing step to say what was found. + if ( + state.step + 1 >= (state.agent.steps ?? MAX_STEPS) || + !hasStepBudget(state.taskBudget) + ) { + const closing = LLM.updateRequest(request, { + messages: [...transcript, Message.user(MAX_STEPS_NOTICE)], + tools: [], + toolChoice: "none", + }) + return Stream.concat( + Stream.fromIterable(results), + runStep(input, tools, withBudget(closing), { + ...state, + step: state.step + 1, + attempt: 0, + closing: true, + }), + ) } - const next = LLM.updateRequest(request, { - messages: [ - ...request.messages, - response.message, - ...dispatched.map(([call, outcome]) => - Message.tool( - ToolResultPart.make({ - id: call.id, - name: call.name, - result: outcome.result, - }), - ), - ), - ], - }) + const next = withBudget(LLM.updateRequest(request, { messages: transcript })) + // A fresh attempt count per step: `attempt` counts retries of *this* step's + // model call, and the shared `budget` is what bounds the turn overall. return Stream.concat( Stream.fromIterable(results), - runStep(input, tools, next, step + 1), + runStep(input, tools, next, { + ...state, + step: state.step + 1, + attempt: 0, + }), ) }), ) diff --git a/apps/api/src/chat/agents.test.ts b/apps/api/src/chat/agents.test.ts new file mode 100644 index 000000000..d84df5212 --- /dev/null +++ b/apps/api/src/chat/agents.test.ts @@ -0,0 +1,99 @@ +/** + * The agent registry's invariants. + * + * These are all "a config mistake becomes a runtime mystery" cases: a mode with no agent, a + * `spawns` entry naming nothing, or a sub-agent whose ruleset asks for an approval it has no way to + * surface. Each is cheap to assert here and expensive to debug in a live turn. + */ +import { ChatMode, makeChatSessionId } from "@maple/domain/chat-session" +import { evaluatePermission } from "@maple/domain/permission" +import { assert, describe, it } from "vitest" +import { AGENTS, agentForSession, buildSystemPrompt, spawnableFor } from "./agents" +import { mapleToolDefinitions } from "@/mcp/tools/registry" + +const subagents = Object.values(AGENTS).filter((agent) => agent.mode === "subagent") + +describe("AGENTS", () => { + it("names an agent for every wire mode", () => { + // `chatModeFromSessionId` is on the wire and the client derives from it, so a mode without a + // matching agent is a runtime `undefined` in the middle of a turn. + for (const mode of ChatMode.literals) { + const agent = AGENTS[mode] + assert.isDefined(agent, `no agent for mode ${mode}`) + assert.equal(agent?.mode, "primary", `${mode} must name a primary agent`) + assert.equal(agent?.name, mode) + } + }) + + it("resolves every spawnable name to a real sub-agent", () => { + for (const agent of Object.values(AGENTS)) { + for (const name of agent.spawns ?? []) { + assert.equal(AGENTS[name]?.mode, "subagent", `${agent.name} spawns unknown ${name}`) + } + } + }) + + it("gives no sub-agent a tool that would need approval", () => { + // A nested turn cannot surface an approval card: approval ends the *outer* turn and is + // applied out of band by `POST /api/chat/apply`. An `ask` here would deadlock the sub-agent + // into proposing something nobody can accept. + for (const agent of subagents) { + for (const definition of mapleToolDefinitions) { + assert.notEqual( + evaluatePermission(agent.permission, definition.name), + "ask", + `${agent.name} would ask for ${definition.name}`, + ) + } + } + }) + + it("denies `task` to every sub-agent, capping nesting structurally", () => { + // The numeric depth cap in the task tool is the belt; this is the braces. A sub-agent whose + // ruleset offered `task` could nest regardless of what the counter said. + for (const agent of subagents) { + assert.equal(evaluatePermission(agent.permission, "task"), "deny", agent.name) + } + }) + + it("has at least one sub-agent, or the delegation machinery is dead code", () => { + assert.isNotEmpty(subagents) + }) +}) + +describe("agentForSession", () => { + it("maps a session id to its mode's agent", () => { + assert.equal(agentForSession(makeChatSessionId("org_1", "tab")).name, "default") + assert.equal( + agentForSession(makeChatSessionId("org_1", "dashboard-builder-123")).name, + "dashboard-builder", + ) + assert.equal(agentForSession(makeChatSessionId("org_1", "inv-abc")).name, "investigate") + }) +}) + +describe("buildSystemPrompt", () => { + it("appends delegation guidance naming exactly the agents this one can spawn", () => { + const prompt = buildSystemPrompt(AGENTS.default!) + + assert.include(prompt, "explore") + assert.include(prompt, "task") + // Generated from the registry, so the prompt and the tool description cannot disagree about + // what is delegable. + for (const agent of subagents.filter((candidate) => candidate.name !== "explore")) { + assert.notInclude(prompt, agent.description) + } + }) + + it("says nothing about delegating when the agent cannot", () => { + const prompt = buildSystemPrompt(AGENTS["dashboard-builder"]!) + + assert.isEmpty(spawnableFor(AGENTS["dashboard-builder"]!)) + assert.notInclude(prompt, "## Delegating") + }) + + it("gives a sub-agent its own persona, not the default one", () => { + assert.notInclude(buildSystemPrompt(AGENTS.explore!), "## Delegating") + assert.include(buildSystemPrompt(AGENTS.explore!), "read-only investigator") + }) +}) diff --git a/apps/api/src/chat/agents.ts b/apps/api/src/chat/agents.ts new file mode 100644 index 000000000..e0231f2d2 --- /dev/null +++ b/apps/api/src/chat/agents.ts @@ -0,0 +1,137 @@ +/** + * The agents a chat turn can run as. + * + * Supersedes `modes.ts`, which was a string switch from `ChatMode` to a system prompt. An agent is + * now a record: a prompt, a permission ruleset, an optional step budget, and the sub-agents it may + * delegate to. That is what makes "this surface should not be able to create alert rules" a + * one-line change instead of a new branch in the loop. + * + * `ChatMode` and `chatModeFromSessionId` are deliberately untouched — they are on the wire and the + * web client derives from them. Every mode names a primary agent, by construction; `agents.test.ts` + * fails if one is ever added without one. + */ +import { chatModeFromSessionId, type ChatMode } from "@maple/domain/chat-session" +import type { PermissionRuleset } from "@maple/domain/permission" +import { DEFAULT_RULESET, READ_ONLY_RULESET } from "./permissions" +import { + DASHBOARD_BUILDER_SYSTEM_PROMPT, + EXPLORE_SYSTEM_PROMPT, + INVESTIGATE_SYSTEM_PROMPT, + SYSTEM_PROMPT, +} from "./prompts" + +export interface AgentDefinition { + readonly name: string + /** Shown to a *calling* model in the `task` tool's description. Written for that reader. */ + readonly description: string + /** `primary` — a session runs as this. `subagent` — only reachable through the `task` tool. */ + readonly mode: "primary" | "subagent" + readonly prompt: string + readonly permission: PermissionRuleset + /** Overrides the turn's default step cap. */ + readonly steps?: number + /** + * Sub-agents this agent may spawn. Empty (the default) means it gets no `task` tool at all — + * the capability is opt-in per agent, not something every turn carries. + */ + readonly spawns?: ReadonlyArray +} + +/** + * Assistant turns a sub-agent gets. + * + * Two-thirds of the parent's budget: enough to search and summarize, not enough to wander. A + * sub-agent that runs out simply reports what it found, which is a fine answer. + */ +export const SUBAGENT_MAX_STEPS = 6 + +export const AGENTS: Readonly> = { + default: { + name: "default", + description: "General Maple assistant.", + mode: "primary", + prompt: SYSTEM_PROMPT, + permission: DEFAULT_RULESET, + spawns: ["explore"], + }, + "dashboard-builder": { + name: "dashboard-builder", + description: "Builds and edits dashboards.", + mode: "primary", + prompt: DASHBOARD_BUILDER_SYSTEM_PROMPT, + permission: DEFAULT_RULESET, + }, + alert: { + name: "alert", + description: "Assists with an alert in context.", + mode: "primary", + prompt: SYSTEM_PROMPT, + permission: DEFAULT_RULESET, + }, + "widget-fix": { + name: "widget-fix", + description: "Repairs a dashboard widget in context.", + mode: "primary", + prompt: SYSTEM_PROMPT, + permission: DEFAULT_RULESET, + }, + investigate: { + name: "investigate", + description: "Runs an autonomous investigation.", + mode: "primary", + prompt: INVESTIGATE_SYSTEM_PROMPT, + permission: DEFAULT_RULESET, + spawns: ["explore"], + }, + explore: { + name: "explore", + description: + "Read-only investigator. Give it a self-contained question about traces, logs, metrics " + + "or errors and it returns a written answer. It cannot change anything and cannot spawn " + + "further agents. Use it to search broadly without filling this conversation with raw " + + "tool output.", + mode: "subagent", + prompt: EXPLORE_SYSTEM_PROMPT, + permission: READ_ONLY_RULESET, + steps: SUBAGENT_MAX_STEPS, + }, +} as const + +/** Every `ChatMode` literal names a primary agent; the mode string *is* the agent name. */ +export const agentForSession = (sessionId: string): AgentDefinition => { + const mode: ChatMode = chatModeFromSessionId(sessionId) + // Non-null by construction, and pinned by `agents.test.ts` rather than by hope. + return AGENTS[mode]! +} + +/** The sub-agents `agent` is allowed to spawn, resolved and filtered to real subagent records. */ +export const spawnableFor = (agent: AgentDefinition): ReadonlyArray => + (agent.spawns ?? []) + .map((name) => AGENTS[name]) + .filter((candidate): candidate is AgentDefinition => candidate?.mode === "subagent") + +/** + * The delegation paragraph appended to a system prompt when an agent can spawn. + * + * The prompt-side twin of the `task` tool's generated description: the tool tells the model *how* + * to call, this tells it *when*. Both are generated from the registry so there is one source of + * truth for what a given agent can delegate to. + */ +const taskGuidance = (spawnable: ReadonlyArray): string => + [ + "## Delegating", + "", + "You can hand a self-contained research question to a sub-agent with the `task` tool. The " + + "sub-agent runs its own tool loop and returns a written answer — its raw tool output never " + + "enters this conversation, so delegation is how you search broadly without burying the " + + "thread in payloads. It sees NOTHING of this conversation, so its prompt must stand alone.", + "", + "Available sub-agents:", + ...spawnable.map((agent) => `- \`${agent.name}\`: ${agent.description}`), + ].join("\n") + +/** The system prompt for a turn: the agent's own persona, plus delegation guidance if it can. */ +export const buildSystemPrompt = (agent: AgentDefinition): string => { + const spawnable = spawnableFor(agent) + return spawnable.length === 0 ? agent.prompt : `${agent.prompt}\n\n${taskGuidance(spawnable)}` +} diff --git a/apps/api/src/chat/context-budget.test.ts b/apps/api/src/chat/context-budget.test.ts new file mode 100644 index 000000000..d45cec708 --- /dev/null +++ b/apps/api/src/chat/context-budget.test.ts @@ -0,0 +1,119 @@ +/** + * The in-turn pruner. + * + * The properties that matter: it protects the recent steps, it never silently shortens a payload + * without saying so, and it does nothing at all when there is nothing worth reclaiming. + */ +import { LLM, Message, ToolResultPart, type LLMRequest, type Model } from "@maple/llm" +import { CloudflareWorkersAI } from "@maple/llm/providers/cloudflare" +import { assert, describe, it } from "vitest" +import { estimateTokens, isNearContextLimit, pruneToolResults } from "./context-budget" + +const MODEL: Model = CloudflareWorkersAI.configure({ accountId: "t", apiKey: "t" }).model("@cf/test/model") + +/** + * Default size clears `PRUNE_MIN_RECLAIM_TOKENS` (20k tokens ≈ 80k chars) from a single pruned + * step, so a test does not have to stack steps just to get past the threshold. + */ +const big = (marker: string, chars = 100_000) => marker + "x".repeat(chars) + +/** One step: an assistant turn plus the tool result it produced. */ +const step = (marker: string, chars?: number): ReadonlyArray => [ + Message.assistant(`calling ${marker}`), + Message.tool(ToolResultPart.make({ id: marker, name: "search_traces", result: big(marker, chars) })), +] + +const requestOf = (messages: ReadonlyArray): LLMRequest => + LLM.request({ model: MODEL, system: "s", messages: [...messages] }) + +const toolTexts = (request: LLMRequest): ReadonlyArray => + request.messages + .filter((message) => message.role === "tool") + .flatMap((message) => message.content) + .map((part) => (part.type === "tool-result" ? String(part.result.value) : "")) + +describe("pruneToolResults", () => { + it("truncates old tool output and leaves the last two steps intact", () => { + const request = requestOf([...step("a"), ...step("b"), ...step("c"), ...step("d")]) + const pruned = pruneToolResults(request) + + const texts = toolTexts(pruned) + assert.lengthOf(texts, 4) + // The two oldest lost their tails; the two the model is still reasoning about did not. + assert.isBelow(texts[0]!.length, 3_000) + assert.isBelow(texts[1]!.length, 3_000) + assert.equal(texts[2], big("c")) + assert.equal(texts[3], big("d")) + }) + + it("tells the model what it dropped instead of quietly shortening the payload", () => { + // A tool result that silently loses its tail is indistinguishable from a tool that returned + // less than it did — which is how a model concludes "there were only 3 traces". + const pruned = pruneToolResults(requestOf([...step("a"), ...step("b"), ...step("c")])) + + assert.include(toolTexts(pruned)[0]!, "truncated") + assert.include(toolTexts(pruned)[0]!, "characters omitted") + }) + + it("keeps the surviving prefix, so the excerpt is the start of the real output", () => { + const pruned = pruneToolResults(requestOf([...step("a"), ...step("b"), ...step("c")])) + + assert.isTrue(toolTexts(pruned)[0]!.startsWith("a")) + }) + + it("returns the request unchanged when there is nothing worth reclaiming", () => { + // Identity, not just equality: a short turn must pay nothing, so callers can apply this + // unconditionally. + const request = requestOf([...step("a", 10), ...step("b", 10), ...step("c", 10)]) + assert.strictEqual(pruneToolResults(request), request) + }) + + it("returns the request unchanged when every step is still protected", () => { + const request = requestOf([...step("a"), ...step("b")]) + assert.strictEqual(pruneToolResults(request), request) + }) + + it("leaves user and assistant messages alone", () => { + const request = requestOf([ + Message.user("why is checkout slow?"), + ...step("a"), + ...step("b"), + ...step("c"), + ]) + const pruned = pruneToolResults(request) + + const spoken = pruned.messages.filter((message) => message.role !== "tool") + assert.deepEqual( + spoken.map((message) => + message.content.map((part) => (part.type === "text" ? part.text : "")).join(""), + ), + ["why is checkout slow?", "calling a", "calling b", "calling c"], + ) + }) +}) + +describe("isNearContextLimit", () => { + it("fires once the input is within the reply's headroom of the window", () => { + assert.isTrue(isNearContextLimit(112_000, { context: 128_000, output: 16_000 })) + assert.isFalse(isNearContextLimit(90_000, { context: 128_000, output: 16_000 })) + }) + + it("caps the reserve so a huge max-output does not make every turn look full", () => { + // Reserve is `min(20k, output)`: a model advertising 128k completion tokens against a + // 1.05M window must not be treated as overflowing at 922k. + assert.isFalse(isNearContextLimit(922_000, { context: 1_050_000, output: 128_000 })) + assert.isTrue(isNearContextLimit(1_040_000, { context: 1_050_000, output: 128_000 })) + }) + + it("says no when the model declares no window, rather than guessing", () => { + assert.isFalse(isNearContextLimit(10_000_000, {})) + }) +}) + +describe("estimateTokens", () => { + it("is a rough chars-per-token estimate, used only to size a prune", () => { + assert.equal(estimateTokens(""), 0) + assert.equal(estimateTokens("abcd"), 1) + assert.equal(estimateTokens("abcde"), 2) + }) +}) diff --git a/apps/api/src/chat/context-budget.ts b/apps/api/src/chat/context-budget.ts new file mode 100644 index 000000000..96ef2cc89 --- /dev/null +++ b/apps/api/src/chat/context-budget.ts @@ -0,0 +1,137 @@ +/** + * Keeping a turn inside the model's context window. + * + * Two different growth problems hide behind "the conversation got too long", and they want + * different fixes: + * + * - **In-turn.** `runStep` rebuilds the request as `[...messages, assistant, ...toolResults]` for + * up to `MAX_STEPS` steps at `TOOL_CONCURRENCY = 4`, and Maple's MCP tools return warehouse rows + * and trace payloads. This is where the tokens actually are, and it is what walks a single + * investigation into the wall. + * - **Cross-turn.** `toLlmMessages` in `turn-runner.ts` replays the transcript, but it already + * drops tool messages entirely and keeps only text, so it is comparatively small. + * + * This module handles the first, and does it without a model call, a wire change, or any + * persistence: old tool output is replaced by a truncated prefix plus an explicit marker. The model + * is *told* what was dropped rather than silently handed a shorter payload, because a tool result + * that quietly loses its tail is indistinguishable from a tool that returned less than it did. + */ +import { LLM, Message, ToolResultPart, type LLMRequest, type ToolResultValue } from "@maple/llm" + +/** + * Rough tokens-per-character. There is no tokenizer in a Worker isolate, and shipping one would + * cost more startup CPU than the estimate is worth. + * + * This is only ever used to decide *whether a prune is worth doing*. The authoritative number — the + * one the trigger reads — is the provider's own `usage.inputTokens` from the previous step. + */ +const CHARS_PER_TOKEN = 4 + +export const tokensFromChars = (chars: number): number => Math.ceil(chars / CHARS_PER_TOKEN) + +export const estimateTokens = (text: string): number => tokensFromChars(text.length) + +/** + * How many of the most recent steps keep their tool output verbatim. + * + * The model is usually still reasoning about the last step or two; truncating those is what makes a + * pruner feel like amnesia rather than housekeeping. + */ +const PROTECT_RECENT_STEPS = 2 + +/** How much of an old tool result survives. */ +const PRUNE_TOOL_RESULT_CHARS = 2_000 + +/** + * Don't bother unless the prune buys real headroom. A turn that would reclaim a few hundred tokens + * pays the rebuild cost and the fidelity cost for nothing. + */ +const PRUNE_MIN_RECLAIM_TOKENS = 20_000 + +/** Headroom kept free for the model's own reply when deciding whether the input still fits. */ +const RESERVE_CEILING_TOKENS = 20_000 + +/** + * Whether an observed input-token count is close enough to the wall to act on. + * + * `undefined` limits mean "we don't know this model's window" — in which case the honest answer is + * no, rather than a guess that could prune a turn that had plenty of room. + */ +export const isNearContextLimit = ( + inputTokens: number, + limits: { readonly context?: number; readonly output?: number }, +): boolean => { + if (limits.context === undefined) return false + const reserved = Math.min(RESERVE_CEILING_TOKENS, limits.output ?? RESERVE_CEILING_TOKENS) + return inputTokens >= limits.context - reserved +} + +const MARKER = (omitted: number) => `\n\n…[truncated, ${omitted} characters omitted]` + +/** Render a tool result as the text the model would see, for measuring and truncating. */ +const resultText = (result: ToolResultValue): string => + typeof result.value === "string" ? result.value : JSON.stringify(result.value) + +/** + * Truncate one tool result, or return `undefined` if it is already short enough. + * + * A string value keeps its original result type; anything else collapses to `"text"`, because a + * truncated JSON payload is no longer parseable and handing the model a broken object is worse than + * handing it a clearly-marked excerpt. + */ +const truncateResult = (result: ToolResultValue): ToolResultValue | undefined => { + const text = resultText(result) + if (text.length <= PRUNE_TOOL_RESULT_CHARS) return undefined + const kept = text.slice(0, PRUNE_TOOL_RESULT_CHARS) + MARKER(text.length - PRUNE_TOOL_RESULT_CHARS) + return typeof result.value === "string" + ? { type: result.type, value: kept } + : { type: "text", value: kept } +} + +/** + * The index of the first message belonging to the last `PROTECT_RECENT_STEPS` steps. + * + * Steps are delimited by assistant messages — `runStep` appends exactly one per step, followed by + * its tool results — so counting assistant turns backwards finds the boundary without the request + * having to carry a step marker. + */ +const protectedFrom = (messages: ReadonlyArray): number => { + let seen = 0 + for (let i = messages.length - 1; i >= 0; i--) { + if (messages[i]?.role !== "assistant") continue + seen += 1 + // The assistant turn that *opens* the oldest protected step. Everything before it belongs + // to an older step and is fair game; the message at this index is not. + if (seen === PROTECT_RECENT_STEPS) return i + } + return 0 +} + +/** + * Truncate tool output from all but the most recent steps. + * + * Returns the request **unchanged** when there is nothing worth reclaiming, so a short turn pays + * nothing and callers can apply this unconditionally. + */ +export const pruneToolResults = (request: LLMRequest): LLMRequest => { + const boundary = protectedFrom(request.messages) + if (boundary === 0) return request + + let reclaimed = 0 + const messages = request.messages.map((message, index) => { + if (index >= boundary || message.role !== "tool") return message + let changed = false + const content = message.content.map((part) => { + if (part.type !== "tool-result") return part + const truncated = truncateResult(part.result) + if (truncated === undefined) return part + changed = true + reclaimed += resultText(part.result).length - resultText(truncated).length + return ToolResultPart.make({ ...part, result: truncated }) + }) + return changed ? Message.make({ ...message, content }) : message + }) + + if (tokensFromChars(reclaimed) < PRUNE_MIN_RECLAIM_TOKENS) return request + return LLM.updateRequest(request, { messages }) +} diff --git a/apps/api/src/chat/llm-retry.test.ts b/apps/api/src/chat/llm-retry.test.ts new file mode 100644 index 000000000..4618f7896 --- /dev/null +++ b/apps/api/src/chat/llm-retry.test.ts @@ -0,0 +1,69 @@ +/** + * The retry classifier. + * + * The load-bearing case is `Transport` / `InvalidProviderOutput` with `retryable: false` — those + * are what a mid-stream failure actually looks like, and a classifier that only trusted the + * `retryable` flag would pass a naive test while doing nothing. See `llm-retry.ts`. + */ +import { LlmCallError } from "@maple/domain/llm" +import { assert, describe, it } from "vitest" +import { isRetryableStepFailure, MAX_STEP_ATTEMPTS, stepRetryDelayMs } from "./llm-retry" + +const error = (patch: { reason: string; retryable?: boolean; contextOverflow?: boolean }): LlmCallError => + new LlmCallError({ + operation: "chat.turn", + module: "test", + method: "stream", + message: "boom", + reason: patch.reason, + retryable: patch.retryable ?? false, + contextOverflow: patch.contextOverflow ?? false, + }) + +describe("isRetryableStepFailure", () => { + it("retries what the provider flags as retryable", () => { + assert.isTrue(isRetryableStepFailure(error({ reason: "RateLimit", retryable: true }))) + assert.isTrue(isRetryableStepFailure(error({ reason: "ProviderInternal", retryable: true }))) + }) + + it("retries stream-level failures despite retryable being false", () => { + // `TransportReason` and `InvalidProviderOutputReason` both hardcode `retryable = false` in + // `lib/llm/src/schema/errors.ts`, and both are how a body that dies mid-stream surfaces. + assert.isTrue(isRetryableStepFailure(error({ reason: "Transport" }))) + assert.isTrue(isRetryableStepFailure(error({ reason: "InvalidProviderOutput" }))) + }) + + it("does not retry failures that will fail identically", () => { + assert.isFalse(isRetryableStepFailure(error({ reason: "Authentication" }))) + assert.isFalse(isRetryableStepFailure(error({ reason: "InvalidRequest" }))) + assert.isFalse(isRetryableStepFailure(error({ reason: "QuotaExceeded" }))) + assert.isFalse(isRetryableStepFailure(error({ reason: "ContentPolicy" }))) + }) + + it("never retries a context overflow, even when the provider says retryable", () => { + // The transcript has to shrink first; the same request cannot start fitting. + assert.isFalse( + isRetryableStepFailure( + error({ reason: "InvalidRequest", retryable: true, contextOverflow: true }), + ), + ) + }) +}) + +describe("stepRetryDelayMs", () => { + it("backs off exponentially and caps", () => { + assert.equal(stepRetryDelayMs(0), 1_000) + assert.equal(stepRetryDelayMs(1), 2_000) + assert.equal(stepRetryDelayMs(2), 4_000) + assert.equal(stepRetryDelayMs(3), 8_000) + assert.equal(stepRetryDelayMs(10), 8_000) + }) + + it("keeps every attempt's backoff under the SSE idle budget", () => { + // `ChatSession.SUBSCRIBE_IDLE_MS` is 25s; a longer gap would recycle the connection + // mid-answer. + for (let attempt = 0; attempt < MAX_STEP_ATTEMPTS; attempt++) { + assert.isBelow(stepRetryDelayMs(attempt), 25_000) + } + }) +}) diff --git a/apps/api/src/chat/llm-retry.ts b/apps/api/src/chat/llm-retry.ts new file mode 100644 index 000000000..86630275e --- /dev/null +++ b/apps/api/src/chat/llm-retry.ts @@ -0,0 +1,77 @@ +/** + * Which model-call failures a chat step retries, and how long it waits. + * + * Pure policy, deliberately separate from the loop that applies it: `apps/api/src/chat/agent.ts` + * retries a *stream*, so it needs the retraction machinery in `runStepAttempt`, while + * `apps/api/src/workflows/triage-agent.ts` retries an `Effect` and can use `Effect.retry` + * directly. The classification must not diverge between them, so it lives here. + * + * ## Why this cannot just read `LlmCallError.retryable` + * + * `lib/llm`'s `RequestExecutor` already retries — twice, capped at 10s — but it wraps + * `executeOnce`, which resolves when *response headers* arrive. It cannot replay a stream that + * died at token 400. That mid-stream case is the entire reason this module exists. + * + * And a mid-stream failure does not report itself as retryable. `route/client.ts`'s stream-level + * `Stream.catchCause` surfaces it as either `Transport` (the fetch body died) or + * `InvalidProviderOutput` (framing or decoding blew up), and both of those reasons hardcode + * `get retryable() { return false }` in `lib/llm/src/schema/errors.ts`. A classifier that trusted + * the flag would retry only what the executor already retried — it would look correct, pass a + * naive test, and do nothing. + */ +import type { LlmCallError } from "@maple/domain/llm" + +/** 1 initial attempt + 3 retries. */ +export const MAX_STEP_ATTEMPTS = 4 + +const STEP_RETRY_BASE_MS = 1_000 +const STEP_RETRY_FACTOR = 2 + +/** + * Ceiling on a single backoff. + * + * Bounded well below `ChatSession`'s `SUBSCRIBE_IDLE_MS` (25s), which recycles an SSE connection + * after that much silence. The `turn-retry` event is itself an append and resets that timer, but a + * longer gap *after* it would still recycle the connection mid-answer. + */ +const STEP_RETRY_MAX_MS = 8_000 + +/** + * Whole-turn ceiling on time spent in backoff, across every step. + * + * opencode retries without an attempt ceiling because it runs in a long-lived process where waiting + * costs nothing but patience. Here the turn holds `ChatSession`'s single turn slot, and + * `TURN_STALE_MS` (15 minutes) will declare it abandoned and append a terminal event *underneath* a + * still-running turn. With `MAX_STEPS = 10`, an uncapped per-step retry multiplies past that. + */ +export const STEP_RETRY_BUDGET_MS = 60_000 + +/** Mutable, shared across a turn's steps — the same pattern as `TurnUsage`, for the same reason. */ +export interface StepRetryBudget { + spentMs: number +} + +export const makeStepRetryBudget = (): StepRetryBudget => ({ spentMs: 0 }) + +/** + * Reasons the vendored executor cannot cover, because they arrive after response headers. + * + * `InvalidProviderOutput` is the debatable inclusion: genuinely malformed provider output will fail + * identically on every attempt and burn the whole budget. It is here because an interrupted body + * also lands on it — `route/client.ts`'s `streamError` only looks for a `Fail` reason and falls + * through to `eventError` otherwise — and `MAX_STEP_ATTEMPTS` bounds the waste to a few seconds. + * Narrow this to `Transport` alone if it proves noisy in practice. + */ +const STREAM_LEVEL_REASONS: ReadonlySet = new Set(["Transport", "InvalidProviderOutput"]) + +export const isRetryableStepFailure = (error: LlmCallError): boolean => { + // Never retried as-is: the request is too big, and sending it again cannot change that. The + // caller shrinks the transcript and retries the *pruned* request instead. + if (error.contextOverflow) return false + // `RateLimit` and `ProviderInternal` — the latter already covers 5xx. + if (error.retryable) return true + return STREAM_LEVEL_REASONS.has(error.reason) +} + +export const stepRetryDelayMs = (attempt: number): number => + Math.min(STEP_RETRY_BASE_MS * STEP_RETRY_FACTOR ** attempt, STEP_RETRY_MAX_MS) diff --git a/apps/api/src/chat/modes.ts b/apps/api/src/chat/modes.ts deleted file mode 100644 index c2c002aa1..000000000 --- a/apps/api/src/chat/modes.ts +++ /dev/null @@ -1,30 +0,0 @@ -// Conversation modes → system prompt. -// -// Moved into apps/api from apps/chat-flue when the chat shell came off Flue. The mode itself lives -// in the shared wire contract (`chatModeFromSessionId` in `@maple/domain/chat-session`), so the -// client, the Durable Object and this prompt assembly all derive it the same way; only the prompt -// text is here. -// -// **Per-conversation context is not assembled here.** Alert / widget-fix / page context reaches the -// model as a fenced preamble on the conversation's first user message, built client-side by -// `apps/web/src/components/chat/context-preamble.ts` and stripped from the visible bubble by -// `stripChatContext`. This file used to carry a second, server-side copy of those same formatters, -// inherited from the pre-Flue agent that received context out-of-band in the request body. Nothing -// reached them — `ChatSendRequest` carries only `text`, so there was no wire path for that context -// to arrive on — and keeping two divergent formatters for one prompt block was worse than keeping -// the one the client actually uses. - -import type { ChatMode } from "@maple/domain/chat-session" -import { DASHBOARD_BUILDER_SYSTEM_PROMPT, INVESTIGATE_SYSTEM_PROMPT, SYSTEM_PROMPT } from "./prompts" - -export interface BuildSystemPromptArgs { - mode: ChatMode -} - -/** The system prompt for a turn, chosen by the conversation's mode. */ -export const buildSystemPrompt = ({ mode }: BuildSystemPromptArgs): string => - mode === "dashboard-builder" - ? DASHBOARD_BUILDER_SYSTEM_PROMPT - : mode === "investigate" - ? INVESTIGATE_SYSTEM_PROMPT - : SYSTEM_PROMPT diff --git a/apps/api/src/chat/permissions.ts b/apps/api/src/chat/permissions.ts new file mode 100644 index 000000000..0670439ba --- /dev/null +++ b/apps/api/src/chat/permissions.ts @@ -0,0 +1,42 @@ +/** + * The rulesets Maple's chat agents run under. + * + * `MUTATING_TOOL_NAMES` stays exactly where it is and keeps its shape: it seeds `DEFAULT_RULESET` + * here, and it remains the allowlist floor for `POST /api/chat/apply`. That is the whole migration + * story — the mirror in `apps/slack-agent/agent/lib/approval.ts` needs no change, and the + * equivalence is pinned by a test in `apps/api/src/mcp/tools/mutating.test.ts` so day-one behaviour + * cannot drift by accident. + */ +import { PermissionRule, type PermissionRuleset } from "@maple/domain/permission" +import { MUTATING_TOOL_NAMES } from "@/mcp/tools/mutating" +import { mapleToolDefinitions } from "@/mcp/tools/registry" + +/** + * Today's behaviour, expressed as data: everything runs, mutations stop and ask. + * + * Sorted so the ruleset is stable across builds — it is a value that will end up in logs and, + * eventually, in a settings UI diff. + */ +export const DEFAULT_RULESET: PermissionRuleset = [ + new PermissionRule({ tool: "*", action: "allow" }), + ...[...MUTATING_TOOL_NAMES].sort().map((tool) => new PermissionRule({ tool, action: "ask" })), +] + +/** + * Read-only: deny everything, then name the tools that are allowed. + * + * Deliberately an allowlist of *concrete registered names* rather than `deny "*"` plus `allow + * "get_*"` globs. A mutating tool added next month is denied by default under this ruleset instead + * of slipping through whatever glob happened to match its name. + * + * Note what this also denies by omission: `task`. A sub-agent running under this ruleset physically + * cannot spawn another one, because `buildChatTools` never offers it the tool. + */ +export const READ_ONLY_RULESET: PermissionRuleset = [ + new PermissionRule({ tool: "*", action: "deny" }), + ...mapleToolDefinitions + .filter((definition) => !MUTATING_TOOL_NAMES.has(definition.name)) + .map((definition) => definition.name) + .sort() + .map((tool) => new PermissionRule({ tool, action: "allow" })), +] diff --git a/apps/api/src/chat/prompts.ts b/apps/api/src/chat/prompts.ts index fd42652de..a4ce05880 100644 --- a/apps/api/src/chat/prompts.ts +++ b/apps/api/src/chat/prompts.ts @@ -4,7 +4,7 @@ // as `mcp__maple__` (e.g. `mcp__maple__find_errors`). The prompts below // keep the short names for readability and add the prefix note once, up front — // the model maps them. Mutating tools follow the propose-then-apply pattern -// (see modes.ts / the approval layer): calling one surfaces an approval step in +// (see agents.ts / the approval layer): calling one surfaces an approval step in // the UI before it takes effect. const TOOL_PREFIX_NOTE = `## Tools @@ -301,3 +301,54 @@ ${APPROVAL_NOTE} - DO NOT narrate your tool calls or explain your investigation process in detail - After adding widgets, confirm what was added in one sentence ` + +/** + * The `explore` sub-agent. + * + * Written for a reader with no conversational context at all: it is handed one self-contained + * question and its final message is the only thing that reaches the parent turn. So the prompt's + * whole job is to make that final message self-sufficient — the raw tool output it looked at is + * discarded, and anything it does not write down is lost. + */ +export const EXPLORE_SYSTEM_PROMPT = `You are a read-only investigator inside the Maple observability platform, working on behalf of another agent. + +${TOOL_PREFIX_NOTE} + +## What you were given +One self-contained question. You cannot see the conversation that produced it, and you cannot ask a follow-up. If the question is ambiguous, investigate the most useful reading of it and say which reading you took. + +## What you can do +Read-only tools only: searching traces, logs, metrics and errors, listing services, and running queries. You cannot create, update or delete anything, and you cannot delegate further. If answering would require a change, say so instead of attempting it. + +## What to return +Your final message is the ONLY thing the caller receives — your tool calls and their output are discarded. Write it so it stands alone: + +- Lead with the answer, not with what you did. +- Include the specific evidence: service and operation names, trace ids, error fingerprints, counts, percentiles, time ranges. These are what the caller needs to act or to drill in, and it cannot get them from you any other way. +- State what you could NOT determine, and why. A confident answer built on a gap is worse than an honest gap. +- No preamble, no offer to help further, no restating of the question. + +Be thorough in your investigation and brief in your report.` + +/** + * The compaction agent. + * + * Its output replaces the head of a long conversation in what the model is replayed. So the bar is + * not "readable summary" — it is "everything a continuation needs, because the originals are gone + * from the model's view". Entity ids matter more than prose here: a summary that says "the checkout + * service was slow" without the trace ids has thrown away the investigation. + */ +export const COMPACTION_SYSTEM_PROMPT = `You are compacting the earlier part of a debugging conversation so it can be carried forward in a smaller context. + +Write a dense factual summary of what happened. Cover: + +- What the user asked for, and any constraints or preferences they stated. +- What was found, with the specific identifiers: service names, operation names, trace ids, error fingerprints, dashboard and alert ids, metric names, time ranges, and the numbers (counts, percentiles, rates). +- What was decided or changed, including anything the user approved or rejected. +- What is still open: unanswered questions, things that were tried and did not work, and anything the user was about to do next. + +Rules: +- Prose and short lists. No headings, no preamble, no sign-off, no offer to help. +- Preserve identifiers verbatim. A summary without them cannot be continued from. +- Do not speculate or add conclusions that were not reached. If something was uncertain, say it was uncertain. +- Write about the conversation in the past tense, as a record. Do not address the user.` diff --git a/apps/api/src/chat/task-tool.test.ts b/apps/api/src/chat/task-tool.test.ts new file mode 100644 index 000000000..98eb2dc1a --- /dev/null +++ b/apps/api/src/chat/task-tool.test.ts @@ -0,0 +1,292 @@ +/** + * Delegating to a sub-agent. + * + * The properties worth pinning are the ones that would be silent if broken: the context firewall + * (the parent must never see the child's tool payloads), the event tagging (a child's events must + * land in the card, not the conversation), and the budgets (a `ToolFailure`, never a defect). + */ +import { LLMClient, type LLMEvent, type LLMRequest, type Model } from "@maple/llm" +import { CloudflareWorkersAI } from "@maple/llm/providers/cloudflare" +import { Effect, Layer, Stream } from "effect" +import { assert, describe, it } from "vitest" +import { runChatTurn, type ChatTurnEvent } from "./agent" +import { AGENTS } from "./agents" +import type { TenantContext } from "@/services/auth/tenant-context" + +const TENANT: TenantContext = { + orgId: "org_test" as TenantContext["orgId"], + userId: "user_test" as TenantContext["userId"], + roles: [], + authMode: "self_hosted", +} + +const MODEL: Model = CloudflareWorkersAI.configure({ accountId: "t", apiKey: "t" }).model("@cf/test/model") + +const textDelta = (text: string): LLMEvent => ({ type: "text-delta", id: "t1", text }) as LLMEvent +const finish = (): LLMEvent => ({ type: "finish", reason: "stop" }) as LLMEvent +const call = (id: string, name: string, input: unknown = {}): LLMEvent => + ({ type: "tool-call", id, name, input, providerExecuted: false }) as LLMEvent + +const taskCall = (id: string, prompt = "why is checkout slow?"): LLMEvent => + call(id, "task", { description: "trace checkout", prompt, subagent_type: "explore" }) + +/** + * Script the model by *which tools the request offers*, not by call order. + * + * The parent and the sub-agent interleave — the child's turn runs inside the parent's tool + * dispatch — so a positional script would depend on scheduling. Keying on the presence of the + * `task` tool identifies the caller unambiguously: only an agent that can spawn is offered it. + */ +const stubModel = ( + parent: ReadonlyArray>, + child: ReadonlyArray>, +) => { + let parentStep = 0 + let childStep = 0 + const service = { + prepare: () => Effect.die(new Error("unused")), + generate: () => Effect.die(new Error("unused")), + stream: (request: LLMRequest) => { + const isParent = request.tools.some((tool) => tool.name === "task") + const scripted = isParent ? parent[parentStep++] : child[childStep++] + return Stream.fromIterable(scripted ?? [finish()]) + }, + } + return { + layer: Layer.succeed(LLMClient.Service, service as never), + childCalls: () => childStep, + } +} + +const run = ( + parent: ReadonlyArray>, + child: ReadonlyArray>, + overrides: { readonly isCurrent?: () => boolean } = {}, +) => { + const stub = stubModel(parent, child) + const emitted: Array = [] + return runChatTurn({ + sessionId: "org_test:tab", + tenant: TENANT, + model: MODEL, + messages: [], + messageId: "m1", + emit: (event) => emitted.push(event), + ...(overrides.isCurrent ? { isCurrent: overrides.isCurrent } : {}), + }).pipe( + Stream.runCollect, + Effect.map((events) => ({ + events: Array.from(events) as ChatTurnEvent[], + emitted, + childCalls: stub.childCalls(), + })), + Effect.provide(stub.layer), + ) +} + +const resultsOf = (events: ReadonlyArray) => + events.filter((event) => event.type === "tool-result") + +describe("the task tool", () => { + it("returns the sub-agent's final text, wrapped, and none of its tool payloads", async () => { + const result = await Effect.runPromise( + run( + [ + [taskCall("t1"), finish()], + [textDelta("Answered."), finish()], + ], + [ + [call("x1", "search_traces"), finish()], + [textDelta("p99 is 4.2s in checkout-api."), finish()], + ], + ), + ) + + const output = String( + resultsOf(result.events)[0]?.type === "tool-result" ? resultsOf(result.events)[0]!.output : "", + ) + assert.include(output, "") + assert.include(output, "p99 is 4.2s in checkout-api.") + // The firewall. If the child's raw output reached the parent, delegation would be strictly + // worse than calling the tools inline — it would cost an extra model call for nothing. + assert.notInclude(output, "search_traces") + }) + + it("drops narration that preceded a tool call, keeping only the report", async () => { + const result = await Effect.runPromise( + run( + [ + [taskCall("t1"), finish()], + [textDelta("ok"), finish()], + ], + [ + [textDelta("Let me check the traces."), call("x1", "search_traces"), finish()], + [textDelta("The answer."), finish()], + ], + ), + ) + + const output = String( + resultsOf(result.events)[0]?.type === "tool-result" ? resultsOf(result.events)[0]!.output : "", + ) + assert.include(output, "The answer.") + assert.notInclude(output, "Let me check") + }) + + it("tags every sub-agent event so it lands in the card, not the conversation", async () => { + const result = await Effect.runPromise( + run( + [ + [taskCall("t1"), finish()], + [textDelta("done"), finish()], + ], + [[textDelta("child text"), finish()]], + ), + ) + + assert.isNotEmpty(result.emitted) + for (const event of result.emitted) { + assert.equal( + event.type === "user-message" ? undefined : event.task?.id, + "t1", + `untagged ${event.type} would surface as a stray top-level message`, + ) + } + // And nothing from the child leaked into the parent's own stream. + const parentText = result.events + .filter((event) => event.type === "text-delta") + .map((event) => (event.type === "text-delta" ? event.text : "")) + .join("") + assert.equal(parentText, "done") + }) + + it("refuses to delegate past the per-turn budget, as a tool failure not a defect", async () => { + // Five task calls against a budget of four. The fifth must come back as a tool result the + // model can route around, not kill the turn. + const result = await Effect.runPromise( + run( + [ + [ + taskCall("t1"), + taskCall("t2"), + taskCall("t3"), + taskCall("t4"), + taskCall("t5"), + finish(), + ], + [textDelta("done"), finish()], + ], + Array.from({ length: 6 }, () => [textDelta("child answer"), finish()]), + ), + ) + + const results = resultsOf(result.events) + assert.lengthOf(results, 5) + const failures = results.filter((event) => event.type === "tool-result" && event.isError === true) + assert.lengthOf(failures, 1) + assert.include(String(failures[0]?.type === "tool-result" ? failures[0].output : ""), "slots left") + // The turn still completed normally. + assert.lengthOf( + result.events.filter((event) => event.type === "turn-end"), + 1, + ) + }, 30_000) + + it("is not offered to an agent with no sub-agents", async () => { + let offered: ReadonlyArray = [] + const service = { + prepare: () => Effect.die(new Error("unused")), + generate: () => Effect.die(new Error("unused")), + stream: (request: LLMRequest) => { + offered = request.tools.map((tool) => tool.name) + return Stream.fromIterable([textDelta("hi"), finish()]) + }, + } + + await Effect.runPromise( + runChatTurn({ + sessionId: "org_test:tab", + tenant: TENANT, + model: MODEL, + messages: [], + messageId: "m1", + agent: AGENTS["dashboard-builder"]!, + }).pipe(Stream.runCollect, Effect.provide(Layer.succeed(LLMClient.Service, service as never))), + ) + + // Delegation is opt-in per agent: `dashboard-builder` has no `spawns`. + assert.notInclude(offered, "task") + }) + + it("is not offered to a sub-agent, capping nesting structurally", async () => { + let offered: ReadonlyArray = [] + const service = { + prepare: () => Effect.die(new Error("unused")), + generate: () => Effect.die(new Error("unused")), + stream: (request: LLMRequest) => { + offered = request.tools.map((tool) => tool.name) + return Stream.fromIterable([textDelta("hi"), finish()]) + }, + } + + await Effect.runPromise( + runChatTurn({ + sessionId: "org_test:tab", + tenant: TENANT, + model: MODEL, + messages: [], + messageId: "c1", + agent: AGENTS.explore!, + depth: 1, + }).pipe(Stream.runCollect, Effect.provide(Layer.succeed(LLMClient.Service, service as never))), + ) + + assert.notInclude(offered, "task") + // And it really is read-only: no mutating tool is even visible. + assert.notInclude(offered, "create_alert_rule") + }) + + it("stops the sub-agent when the parent turn is superseded", async () => { + const result = await Effect.runPromise( + run([[taskCall("t1"), finish()]], [[textDelta("child"), finish()]], { + isCurrent: () => false, + }), + ) + + // The child inherits the parent's `isCurrent` closure, so an abort reaches both without any + // separate cancellation path. + assert.equal(result.childCalls, 0) + assert.isEmpty(result.events.filter((event) => event.type === "turn-end")) + }) +}) + +describe("the task tool's description", () => { + it("names exactly the sub-agents the caller may spawn", async () => { + let description: string | undefined + const service = { + prepare: () => Effect.die(new Error("unused")), + generate: () => Effect.die(new Error("unused")), + stream: (request: LLMRequest) => { + description = request.tools.find((tool) => tool.name === "task")?.description + return Stream.fromIterable([textDelta("hi"), finish()]) + }, + } + + await Effect.runPromise( + runChatTurn({ + sessionId: "org_test:tab", + tenant: TENANT, + model: MODEL, + messages: [], + messageId: "m1", + }).pipe(Stream.runCollect, Effect.provide(Layer.succeed(LLMClient.Service, service as never))), + ) + + // Generated from the registry, so the tool description and the system prompt cannot disagree + // about what is delegable. + assert.include(description ?? "", "explore") + assert.include(description ?? "", AGENTS.explore!.description) + // The one thing the model most needs to know, and the one it will otherwise get wrong. + assert.include(description ?? "", "sees NOTHING of this conversation") + }) +}) diff --git a/apps/api/src/chat/task-tool.ts b/apps/api/src/chat/task-tool.ts new file mode 100644 index 000000000..dc66ba4a4 --- /dev/null +++ b/apps/api/src/chat/task-tool.ts @@ -0,0 +1,234 @@ +/** + * `task` — delegating a self-contained question to a sub-agent. + * + * ## Why this runs in process, and not in a Durable Object of its own + * + * opencode's equivalent creates a real child *session* and re-enters the same loop against it. That + * is nearly free there, because a session is an in-memory record in a long-lived process. On + * Cloudflare it inverts: a child `ChatSession` DO would need a new addressing scheme that survives + * all four session-id parsers (`orgIdFromChatSessionId`, `tabIdFromChatSessionId`, + * `chatModeFromSessionId`, `investigationIdFromChatSessionId`), a distributed abort fan-out that + * survives isolate eviction, and either cross-DO event mirroring or a second SSE stream per task — + * all to solve a CPU-headroom problem the budgets below already bound. + * + * So the sub-agent runs *inside* the parent turn: `runChatTurn` calls itself, with a different + * agent record and its own step budget. Three things fall out for free that the DO design would + * have had to build: + * + * - **Abort.** The child gets the parent's `isCurrent` closure verbatim, so clearing the turn + * claim stops both at their next step boundary. + * - **One log, one stream.** Child events are tagged with a {@link ChatTaskRef} and land in the + * parent's SQLite log, where `history()` folds them into a nested transcript. + * - **Replay.** `toLlmMessages` filters on top-level text, and child text lives nested inside a + * tool call, so the sub-agent's chatter never re-enters the next turn's context. + * + * ## The context firewall + * + * The parent gets back the child's **final assistant text only**, never its tool payloads. That is + * the entire point: delegation is how the model searches broadly without burying the conversation + * in warehouse rows. A `task` that returned the child's transcript would be strictly worse than + * just calling the tools inline. + */ +import { Schema, Semaphore, Stream, Effect } from "effect" +import { Tool, ToolFailure, type Tools } from "@maple/llm" +import { Message } from "@maple/llm" +import { AGENTS, type AgentDefinition } from "./agents" +import type { ChatTurnEvent, ChatTurnInput } from "./agent" + +/** + * How deep nesting may go. 1 means the primary turn may spawn, and sub-agents may not. + * + * Capped twice, as opencode does: structurally, because every sub-agent's ruleset denies `task` so + * `buildChatTools` never offers it; and numerically, here. Belt and braces, because the structural + * cap silently disappears the moment someone writes a sub-agent ruleset with `{"*": allow}`. + */ +const SUBAGENT_MAX_DEPTH = 1 + +/** Concurrent sub-agent turns inside one Durable Object. */ +const TASK_CONCURRENCY = 2 + +/** Total sub-agent turns one parent turn may start, across all its steps. */ +const TASK_BUDGET_PER_TURN = 4 + +/** + * Assistant turns across the parent and every descendant. + * + * A backstop against the multiplication: 10 parent steps × 4 tasks × 6 child steps is 240 model + * calls and easily tens of minutes, which would reach `ChatSession`'s 15-minute `TURN_STALE_MS` + * watchdog and have it declare the turn abandoned *underneath* a still-running turn. If real turns + * start approaching that, lower `TASK_BUDGET_PER_TURN` — do not raise the watchdog. + */ +const TURN_STEP_BUDGET = 30 + +/** + * Per-turn state shared with every descendant. + * + * Mutable and shared rather than returned, the same shape and for the same reason as `TurnUsage`: + * the consumer is a tool invoked mid-turn, so there is no "after the turn" moment to reconcile at. + */ +export interface TaskBudget { + tasksStarted: number + stepsUsed: number + readonly semaphore: Semaphore.Semaphore +} + +export const makeTaskBudget = (): TaskBudget => ({ + tasksStarted: 0, + stepsUsed: 0, + semaphore: Semaphore.makeUnsafe(TASK_CONCURRENCY), +}) + +/** Whether the turn may start another step at all. Checked by the loop, not just by this tool. */ +export const hasStepBudget = (budget: TaskBudget | undefined): boolean => + budget === undefined || budget.stepsUsed < TURN_STEP_BUDGET + +const renderResult = (ref: { id: string; agent: string }, text: string): string => + `\n\n${text}\n\n` + +const describeSpawnable = (spawnable: ReadonlyArray): string => + [ + "Hand a self-contained research question to a sub-agent. It runs its own tool loop and " + + "returns a written answer; its raw tool output never enters this conversation, so this is " + + "how you search broadly without burying the thread in payloads.", + "", + "The sub-agent sees NOTHING of this conversation — no history, no context, no attachments. " + + "Its `prompt` must stand alone and name everything it needs. You cannot ask it a " + + "follow-up, so ask for everything you need in one go.", + "", + "Launch several concurrently when the questions are independent.", + "", + "Available sub-agents:", + ...spawnable.map((agent) => `- ${agent.name}: ${agent.description}`), + ].join("\n") + +/** + * Build the `task` tool for `agent`, or nothing when it has no sub-agents to spawn. + * + * `runTurn` is injected rather than imported so this module does not import `agent.ts` while + * `agent.ts` imports it — the recursion stays visible at the one call site that wires it. + */ +export const buildTaskTool = ( + input: ChatTurnInput, + spawnable: ReadonlyArray, + budget: TaskBudget, + runTurn: (child: ChatTurnInput) => Stream.Stream, +): Tools => { + if (spawnable.length === 0) return {} + if ((input.depth ?? 0) >= SUBAGENT_MAX_DEPTH) return {} + + const names = spawnable.map((agent) => agent.name) + + return { + task: Tool.make({ + description: describeSpawnable(spawnable), + parameters: Schema.Struct({ + description: Schema.String.annotate({ + description: 'A 3-5 word label for the UI, e.g. "trace checkout latency".', + }), + prompt: Schema.String.annotate({ + description: + "The complete instruction for the sub-agent. It sees nothing of this " + + "conversation, so this must stand alone.", + }), + subagent_type: Schema.Literals(names).annotate({ + description: "Which sub-agent to run.", + }), + }), + success: Schema.String, + execute: (params, context) => + budget.semaphore.withPermits(1)(runTask(input, params, context?.id, budget, runTurn)), + }), + } +} + +const runTask = ( + input: ChatTurnInput, + params: { description: string; prompt: string; subagent_type: string }, + callId: string | undefined, + budget: TaskBudget, + runTurn: (child: ChatTurnInput) => Stream.Stream, +): Effect.Effect => + Effect.gen(function* () { + // Over budget is a `ToolFailure`, never a defect: the model is told to do the work inline + // rather than having the whole turn die because it delegated once too often. + if (budget.tasksStarted >= TASK_BUDGET_PER_TURN) { + return yield* Effect.fail( + new ToolFailure({ + message: + `No sub-agent slots left this turn (limit ${TASK_BUDGET_PER_TURN}). ` + + "Do this part of the work yourself with the tools you have.", + }), + ) + } + if (!hasStepBudget(budget)) { + return yield* Effect.fail( + new ToolFailure({ message: "This turn is out of steps. Answer from what you have." }), + ) + } + + const child = AGENTS[params.subagent_type] + if (child === undefined || child.mode !== "subagent") { + return yield* Effect.fail( + new ToolFailure({ message: `Unknown sub-agent: ${params.subagent_type}` }), + ) + } + + // The call id binds the child's events to the card the parent already announced. There is + // always one in practice; the fallback keeps the tool usable outside a dispatch loop. + const taskId = callId ?? crypto.randomUUID() + budget.tasksStarted += 1 + + const ref = { + id: taskId, + agent: child.name, + parentMessageId: input.messageId, + } as const + + /** + * The child's answer, as the parent will see it. + * + * Reset on every tool call because text a sub-agent emits *before* calling a tool is + * narration ("let me check the traces"), not findings. Only the run of text after its last + * tool call is the report. + */ + let answer = "" + + yield* runTurn({ + sessionId: input.sessionId, + tenant: input.tenant, + model: input.model, + agent: child, + messages: [Message.user(params.prompt)], + messageId: crypto.randomUUID(), + task: ref, + depth: (input.depth ?? 0) + 1, + taskBudget: budget, + // Usage rolls up: a delegated search is still this turn's spend, and the org is billed + // for it either way. + ...(input.usage ? { usage: input.usage } : {}), + ...(input.isCurrent ? { isCurrent: input.isCurrent } : {}), + ...(input.emit ? { emit: input.emit } : {}), + }).pipe( + Stream.runForEach((event) => + Effect.sync(() => { + input.emit?.(event) + if (event.type === "text-delta") answer += event.text + if (event.type === "tool-call") answer = "" + }), + ), + Effect.catchCause(() => + Effect.fail(new ToolFailure({ message: `The ${child.name} sub-agent failed.` })), + ), + ) + + const text = answer.trim() + if (text === "") { + return yield* Effect.fail( + new ToolFailure({ message: `The ${child.name} sub-agent returned no result.` }), + ) + } + + // Wrapped rather than bare so the model reliably reads it as a report from elsewhere instead + // of mistaking it for its own prose. + return renderResult(ref, text) + }) diff --git a/apps/api/src/chat/turn-runner.test.ts b/apps/api/src/chat/turn-runner.test.ts new file mode 100644 index 000000000..de13ed893 --- /dev/null +++ b/apps/api/src/chat/turn-runner.test.ts @@ -0,0 +1,121 @@ +/** + * `toLlmMessages` — what a new turn actually replays to the model. + * + * This is where a long conversation either keeps its beginning or loses it. The head-drop is the + * fallback and must stay byte-for-byte what it was; the compaction path is the improvement. + */ +import type { ChatMessage } from "@maple/domain/chat-session" +import { assert, describe, it } from "vitest" +import { toLlmMessages } from "./turn-runner" + +let seq = 0 + +const message = (role: "user" | "assistant", text: string, toolCalls: unknown[] = []): ChatMessage => + ({ + id: `m${(seq += 1)}`, + role, + text, + toolCalls, + createdAt: seq, + startSeq: seq, + }) as unknown as ChatMessage + +const textOf = (messages: ReadonlyArray<{ content: ReadonlyArray }>) => + messages.map((m) => m.content.map((part) => (part as { text?: string }).text ?? "").join("")) + +describe("toLlmMessages without a compaction", () => { + it("replays the conversation in order", () => { + seq = 0 + const replayed = toLlmMessages([ + message("user", "why is checkout slow?"), + message("assistant", "Looking."), + message("user", "and payments?"), + ]) + + assert.deepEqual(textOf(replayed), ["why is checkout slow?", "Looking.", "and payments?"]) + assert.deepEqual( + replayed.map((m) => m.role), + ["user", "assistant", "user"], + ) + }) + + it("drops messages with no text, so a pure tool turn is not replayed as an empty one", () => { + seq = 0 + const replayed = toLlmMessages([ + message("user", "check it"), + message("assistant", "", [{ id: "c1" }]), + message("assistant", "Done."), + ]) + + assert.deepEqual(textOf(replayed), ["check it", "Done."]) + }) + + it("drops from the head when the transcript exceeds the message cap", () => { + seq = 0 + const history = Array.from({ length: 50 }, (_, i) => message("user", `turn ${i}`)) + const replayed = toLlmMessages(history) + + assert.lengthOf(replayed, 40) + // The tail is what the next turn needs; the head is what a human skimming would skip. + assert.equal(textOf(replayed)[39], "turn 49") + }) + + it("keeps one message even when it alone blows the character budget", () => { + // The `kept.length > 0` guard: otherwise a single pasted stack trace replays as *nothing*, + // and the model answers the next question with no context at all. + seq = 0 + const replayed = toLlmMessages([message("user", "x".repeat(80_000))]) + + assert.lengthOf(replayed, 1) + }) +}) + +describe("toLlmMessages with a compaction", () => { + it("replaces the head with its summary instead of dropping it", () => { + seq = 0 + const history = [ + message("user", "why is checkout slow?"), + message("assistant", "p99 is 4.2s, trace abc123."), + message("user", "and payments?"), + message("assistant", "Payments is fine."), + ] + const replayed = toLlmMessages(history, { + summary: "The user asked about checkout; p99 was 4.2s in trace abc123.", + throughSeq: history[1]!.startSeq, + }) + + const texts = textOf(replayed) + assert.lengthOf(replayed, 3) + assert.include(texts[0], "Summary of the earlier part") + assert.include(texts[0], "trace abc123") + // Everything at or before the split point is gone; everything after is verbatim. + assert.deepEqual(texts.slice(1), ["and payments?", "Payments is fine."]) + }) + + it("carries the summary as a user turn, not a fabricated assistant one", () => { + // A synthetic assistant message would assert the model said something it did not, and + // several protocols dislike a conversation opening on an assistant turn. + seq = 0 + const history = [message("user", "old"), message("user", "new")] + const replayed = toLlmMessages(history, { summary: "s", throughSeq: history[0]!.startSeq }) + + assert.equal(replayed[0]?.role, "user") + }) + + it("still bounds the tail, because a compaction can be arbitrarily stale", () => { + seq = 0 + const history = Array.from({ length: 60 }, (_, i) => message("user", `turn ${i}`)) + const replayed = toLlmMessages(history, { summary: "s", throughSeq: 0 }) + + // One summary plus the capped tail — not sixty-one messages. + assert.lengthOf(replayed, 41) + }) + + it("replays everything when the compaction predates the whole transcript", () => { + seq = 0 + const history = [message("user", "a"), message("user", "b")] + const replayed = toLlmMessages(history, { summary: "s", throughSeq: 0 }) + + assert.deepEqual(textOf(replayed).slice(1), ["a", "b"]) + }) +}) diff --git a/apps/api/src/chat/turn-runner.ts b/apps/api/src/chat/turn-runner.ts index 95bd17f58..152303f67 100644 --- a/apps/api/src/chat/turn-runner.ts +++ b/apps/api/src/chat/turn-runner.ts @@ -26,7 +26,7 @@ import { type ChatTurnTenantEncoded, } from "@maple/domain/chat-session" import { layerFromEnvRecord, WorkerConfigProviderLayer } from "@maple/effect-cloudflare" -import { Message } from "@maple/llm" +import { LLM, Message, type Model } from "@maple/llm" import { Effect, Layer, ManagedRuntime, Stream } from "effect" import type { ChatSession } from "./ChatSession" import type { TenantContext } from "@/services/auth/tenant-context" @@ -76,15 +76,35 @@ const toTenantContext = (encoded: ChatTurnTenantEncoded): TenantContext => { const MAX_REPLAYED_MESSAGES = 40 const MAX_REPLAYED_CHARS = 60_000 +/** How the summary is introduced to the model. */ +const COMPACTION_PREAMBLE = + "Summary of the earlier part of this conversation, which is no longer shown in full:\n\n" + /** * Project the durable transcript into `@maple/llm` messages, most recent first-limited. * * Tool calls are deliberately NOT replayed as tool messages: a rehydrated conversation needs the * *conclusions*, not a second copy of every tool payload, and replaying tool results without their * matching provider-native call ids is what makes providers reject a continuation. The assistant's - * text is what carries forward. + * text is what carries forward. Sub-agent text is likewise excluded for free — it lives nested + * inside a tool call, not at the top level. + * + * With a compaction, the head is replaced by its summary instead of being dropped. Without one, the + * head-drop below is the fallback and stays exactly as it was: a compaction can fail or be evicted + * before it is written, and a degraded turn is much better than a broken one. + * + * The summary rides as a **user** message. A synthetic assistant turn would assert the model said + * something it did not, and several protocols dislike a conversation opening on an assistant turn. */ -const toLlmMessages = (history: ReadonlyArray): ReadonlyArray => { +export const toLlmMessages = ( + history: ReadonlyArray, + compaction?: { readonly summary: string; readonly throughSeq: number }, +): ReadonlyArray => { + if (compaction !== undefined) { + const tail = toLlmMessages(history.filter((message) => message.startSeq > compaction.throughSeq)) + return [Message.user(COMPACTION_PREAMBLE + compaction.summary), ...tail] + } + const spoken = history.filter((message) => message.text.trim() !== "") // Walk backwards so the newest turns are the ones kept: the tail is the part the next turn @@ -105,6 +125,67 @@ const toLlmMessages = (history: ReadonlyArray): ReadonlyArray => + Effect.gen(function* () { + const { contextLimitOf, outputLimitOf } = yield* Effect.promise(() => import("../platform/Llm")) + const { isNearContextLimit } = yield* Effect.promise(() => import("./context-budget")) + if ( + !isNearContextLimit(usage.input, { + context: contextLimitOf(model), + output: outputLimitOf(model), + }) + ) { + return + } + // An aborted turn must not append a compaction: the conversation it would summarize is not + // the one the user is looking at. + if (!input.session.holdsTurn(input.messageId)) return + + const { COMPACTION_SYSTEM_PROMPT } = yield* Effect.promise(() => import("./prompts")) + const history = input.session.history() + const throughSeq = input.session.cursor() + const previous = input.session.compaction() + + const response = yield* LLM.generate( + LLM.request({ + model, + system: COMPACTION_SYSTEM_PROMPT, + messages: [...toLlmMessages(history, previous)], + prompt: "Compact the conversation above now.", + }), + ) + const summary = response.text.trim() + if (summary === "") return + input.session.append({ type: "compaction", messageId: input.messageId, summary, throughSeq }) + }).pipe( + // Bounded, and never allowed to turn a delivered answer into a failed turn. A conversation + // that stays uncompacted just falls back to the head-drop next time. + Effect.timeout(COMPACTION_TIMEOUT), + Effect.catchCause(() => Effect.void), + ) + +/** Compaction is housekeeping; it must not hold the turn slot open. */ +const COMPACTION_TIMEOUT = "20 seconds" + /** * Drive one turn to completion. * @@ -157,18 +238,27 @@ export const runChatSessionTurn = async (input: RunChatSessionTurnInput): Promis sessionId: input.sessionId, tenant, model, - messages: toLlmMessages(history), + messages: toLlmMessages(history, input.session.compaction()), messageId: input.messageId, extraTools, usage, // An abort clears the claim; the turn notices here and stops at the next event // rather than streaming into a conversation that has moved on. isCurrent: () => input.session.holdsTurn(input.messageId), + // Sub-agent events reach the log through here rather than through the returned + // stream — see the note on `ChatTurnInput.emit`. The `holdsTurn` guard mirrors the + // `Stream.takeWhile` below, so an aborted turn's in-flight sub-agent stops writing + // too rather than appending into a conversation that has moved on. + emit: (event) => { + if (input.session.holdsTurn(input.messageId)) input.session.append(event) + }, }) .pipe( Stream.takeWhile(() => input.session.holdsTurn(input.messageId)), Stream.runForEach((event) => Effect.sync(() => input.session.append(event))), ) + + yield* compactIfNeeded(input, model, usage) }).pipe( Effect.withSpan("chat.turn", { attributes: { diff --git a/apps/api/src/mcp/tools/llm-tools.ts b/apps/api/src/mcp/tools/llm-tools.ts new file mode 100644 index 000000000..cc83b99e6 --- /dev/null +++ b/apps/api/src/mcp/tools/llm-tools.ts @@ -0,0 +1,133 @@ +/** + * Maple's MCP registry, wrapped as `@maple/llm` tools. + * + * One wrapper, two callers: the streaming chat turn (`apps/api/src/chat/agent.ts`) and the + * autonomous triage loop (`apps/api/src/workflows/triage-agent.ts`). They were near-identical + * copies, down to a duplicated `withRuntimeServices` and `toolResultText` — and they had already + * drifted in the one place it mattered: triage still rendered a whole Effect `Cause` into the + * message it handed the model, which is stack frames and, inside a `DatabaseError`, connection + * details. Both now go through `summarizeToolFailure`. + * + * Dynamic (`jsonSchema`) mode is the right fit for every tool here: `toInputSchema` already produces + * the exact JSON Schema the MCP surface publishes, and `callMcpTool` does its own Effect Schema + * decode with the tool's own error messages. Decoding twice would only give the model a second, + * worse phrasing of the same validation failure. + * + * The tenant is provided per call rather than ambiently so a loop can never widen its own scope: + * every tool executes under exactly the org the turn was started for. + */ +import { Tool, ToolFailure, type Tools } from "@maple/llm" +import { Cause, Effect } from "effect" +import { callMcpTool } from "@/mcp/dispatcher" +import { CurrentMcpTenant } from "@/mcp/lib/query-warehouse" +import { mapleToolDefinitions, toInputSchema } from "@/mcp/tools/registry" +import type { TenantContext } from "@/services/auth/tenant-context" + +/** + * Re-pin the service requirements of an MCP tool handler. + * + * `MapleToolDefinition.handler` types its requirements as `any` — a deliberate erasure at the + * boundary between ~57 heterogeneous tool implementations and the single `McpServer.addTool` + * signature (see the comment on the type in `mcp/tools/registry.ts`). `@maple/llm`'s `Tool.make` + * insists on `never`, and `any` is not assignable to `never` in that position, so the erasure has + * to be undone somewhere. Doing it here, once and by name, keeps it visible: the services really + * are supplied, by the `ManagedRuntime` each caller builds around its loop. + */ +// eslint-disable-next-line @typescript-eslint/no-explicit-any +export const withRuntimeServices = (effect: Effect.Effect): Effect.Effect => + effect as Effect.Effect + +/** + * Serialize an MCP tool result for the model. Maple's tools already return model-facing text + * blocks, so this is a join rather than a re-encode; `isError` is surfaced as a `ToolFailure` so + * the runtime emits a `tool-error` event and the model can self-correct. + */ +export const toolResultText = (result: { content: ReadonlyArray<{ text: string }> }): string => + result.content.map((block) => block.text).join("\n") + +/** + * A one-line reason for a failed tool, safe to hand the model. + * + * `String(cause)` renders the whole Effect cause: stack frames, and — inside a `DatabaseError` — + * connection details. That went into the model's context and, through the tool-result event, into + * a durable transcript the browser reads back. + */ +export const summarizeToolFailure = (cause: Cause.Cause): string => { + const failure = cause.reasons.find(Cause.isFailReason) + const error: unknown = failure?.error + if (error instanceof Error) return error.message + if (error && typeof error === "object" && "message" in error) { + const message = (error as { message?: unknown }).message + if (typeof message === "string") return message + } + return "the tool failed" +} + +/** + * Description suffix on gated tools. The model still calls them normally; it just needs to know + * the call is a proposal so it stops rather than narrating a completed change. + */ +export const APPROVAL_NOTE = + "\n\nThis is an approval-gated action. Calling it proposes the change for the user to approve; " + + "it does NOT take effect until they do. Call it once with the intended arguments and stop." + +export interface BuildMapleToolsOptions { + /** Which registry tools to expose. Defaults to all of them. */ + readonly include?: (name: string) => boolean + /** + * Which exposed tools are approval-gated. A gated tool keeps a real handler so the schema the + * model sees is identical to the ungated case, but the handler refuses: the caller's loop is + * expected to break on the proposal before ever dispatching it. + */ + readonly gate?: (name: string) => boolean +} + +/** Wrap the Maple MCP registry as `@maple/llm` tools. */ +export const buildMapleTools = (tenant: TenantContext, options: BuildMapleToolsOptions = {}): Tools => + Object.fromEntries( + mapleToolDefinitions + .filter((definition) => options.include?.(definition.name) ?? true) + .map((definition) => { + const gated = options.gate?.(definition.name) ?? false + return [ + definition.name, + Tool.make({ + description: gated + ? `${definition.description}${APPROVAL_NOTE}` + : definition.description, + jsonSchema: toInputSchema(definition.schema), + execute: (params): Effect.Effect => + gated + ? Effect.fail( + new ToolFailure({ + message: `${definition.name} requires user approval and was not executed.`, + }), + ) + : withRuntimeServices( + callMcpTool(definition.name, params).pipe( + Effect.provideService(CurrentMcpTenant, tenant), + Effect.flatMap((result) => + result.isError + ? Effect.fail( + new ToolFailure({ + message: toolResultText(result), + }), + ) + : Effect.succeed(toolResultText(result)), + ), + // A tool that fails outright (unknown tool, tenant error) must + // not kill the turn — hand the model the message and let it + // route around. + Effect.catchCause((cause) => + Effect.fail( + new ToolFailure({ + message: `Tool failed: ${summarizeToolFailure(cause)}`, + }), + ), + ), + ), + ), + }), + ] + }), + ) diff --git a/apps/api/src/mcp/tools/mutating.test.ts b/apps/api/src/mcp/tools/mutating.test.ts index cb2d6e700..db2e7421e 100644 --- a/apps/api/src/mcp/tools/mutating.test.ts +++ b/apps/api/src/mcp/tools/mutating.test.ts @@ -1,6 +1,8 @@ import { describe, expect, it } from "vitest" import { mapleToolDefinitions } from "./registry" import { MUTATING_TOOL_NAMES } from "./mutating" +import { evaluatePermission, isToolVisible } from "@maple/domain/permission" +import { DEFAULT_RULESET, READ_ONLY_RULESET } from "@/chat/permissions" describe("MUTATING_TOOL_NAMES", () => { it("every approval-gated tool exists in the registry", () => { @@ -23,3 +25,52 @@ describe("MUTATING_TOOL_NAMES", () => { expect(MUTATING_TOOL_NAMES.has("transition_error_issue")).toBe(true) }) }) + +describe("DEFAULT_RULESET", () => { + it("is exactly MUTATING_TOOL_NAMES, expressed as rules", () => { + // The migration lock. Rulesets replaced a flat `Set`, and the whole claim of that change is + // that day-one behaviour is unchanged: every registered tool resolves to `ask` if and only + // if it is in the set, and to `allow` otherwise. If either side drifts, this fails loudly + // rather than quietly widening what the chat agent can do without approval. + for (const definition of mapleToolDefinitions) { + expect( + evaluatePermission(DEFAULT_RULESET, definition.name), + `permission drifted for ${definition.name}`, + ).toBe(MUTATING_TOOL_NAMES.has(definition.name) ? "ask" : "allow") + } + }) + + it("hides nothing — the gate is approval, not invisibility", () => { + for (const definition of mapleToolDefinitions) { + expect(isToolVisible(DEFAULT_RULESET, definition.name)).toBe(true) + } + }) +}) + +describe("READ_ONLY_RULESET", () => { + it("denies every mutating tool, so a sub-agent cannot even see one", () => { + for (const name of MUTATING_TOOL_NAMES) { + expect(evaluatePermission(READ_ONLY_RULESET, name), `${name} was visible`).toBe("deny") + } + }) + + it("allows the read-only tools an investigator actually needs", () => { + for (const name of ["find_errors", "search_traces", "list_services", "query_data"]) { + expect(evaluatePermission(READ_ONLY_RULESET, name), `${name} was denied`).toBe("allow") + } + }) + + it("denies a tool that does not exist yet, rather than matching it by glob", () => { + // An allowlist of concrete names, not `deny "*"` plus `allow "get_*"`: a mutating tool added + // next month is denied by default instead of slipping through whatever glob fits its name. + expect(evaluatePermission(READ_ONLY_RULESET, "get_and_then_delete_everything")).toBe("deny") + }) + + it("never asks — a sub-agent turn has no way to surface an approval card", () => { + // Approval ends the *outer* turn and is applied by `POST /api/chat/apply`; a nested turn has + // no such exit, so an `ask` in a sub-agent ruleset is a configuration error. + for (const definition of mapleToolDefinitions) { + expect(evaluatePermission(READ_ONLY_RULESET, definition.name)).not.toBe("ask") + } + }) +}) diff --git a/apps/api/src/platform/Llm.test.ts b/apps/api/src/platform/Llm.test.ts index 39c5739ae..a94ce09e4 100644 --- a/apps/api/src/platform/Llm.test.ts +++ b/apps/api/src/platform/Llm.test.ts @@ -13,7 +13,14 @@ import { LLM } from "@maple/llm" import { Effect, Layer } from "effect" import { FetchHttpClient } from "effect/unstable/http" import { describe, expect, it } from "vitest" -import { layerLlm, resolveTriageModel, type LlmCallTags, type LlmEnv } from "./Llm" +import { + contextLimitOf, + layerLlm, + outputLimitOf, + resolveTriageModel, + type LlmCallTags, + type LlmEnv, +} from "./Llm" interface CapturedRequest { readonly url: string @@ -113,3 +120,43 @@ describe("resolveTriageModel — OpenRouter attribution", () => { expect(captured.body).not.toHaveProperty("trace") }) }) + +describe("resolveTriageModel — context limits", () => { + it("attaches the configured model's window, which the vendored provider leaves unset", () => { + // `@maple/llm` declares `ModelLimits` but no provider populates it, so before this every + // model reported `undefined` and nothing could tell when a transcript was near the wall. + const model = resolveTriageModel(openRouterEnv) + + expect(contextLimitOf(model)).toBe(1_050_000) + expect(outputLimitOf(model)).toBe(128_000) + }) + + it("falls back to a conservative window for a model it does not know", () => { + // Too low costs a summarization call; too high costs the whole turn. Unknown means low. + const model = resolveTriageModel({ + ...openRouterEnv, + MAPLE_TRIAGE_MODEL_OPENROUTER: "some/model-shipped-after-this-table", + }) + + expect(contextLimitOf(model)).toBe(128_000) + }) + + it("lets the environment override the table", () => { + const model = resolveTriageModel({ + ...openRouterEnv, + MAPLE_TRIAGE_MODEL_CONTEXT: "64000", + MAPLE_TRIAGE_MODEL_OUTPUT: "4000", + }) + + expect(contextLimitOf(model)).toBe(64_000) + expect(outputLimitOf(model)).toBe(4_000) + }) + + it("ignores an unparseable or nonsensical override rather than trusting it", () => { + // A zero or negative window would make every turn look overflowed on its first step. + for (const bad of ["", "not-a-number", "0", "-5", "1.5"]) { + const model = resolveTriageModel({ ...openRouterEnv, MAPLE_TRIAGE_MODEL_CONTEXT: bad }) + expect(contextLimitOf(model)).toBe(1_050_000) + } + }) +}) diff --git a/apps/api/src/platform/Llm.ts b/apps/api/src/platform/Llm.ts index ec62784df..10dabebe0 100644 --- a/apps/api/src/platform/Llm.ts +++ b/apps/api/src/platform/Llm.ts @@ -20,8 +20,8 @@ import { LlmCallError } from "@maple/domain/llm" import { CloudflareWorkersAI } from "@maple/llm/providers/cloudflare" import * as OpenRouter from "@maple/llm/providers/openrouter" import { LLMClient, RequestExecutor } from "@maple/llm/route" -import { isContextOverflowFailure } from "@maple/llm" -import type { LLMClientService, LLMError, Model } from "@maple/llm" +import { isContextOverflowFailure, Model } from "@maple/llm" +import type { LLMClientService, LLMError } from "@maple/llm" import { Layer } from "effect" import { FetchHttpClient } from "effect/unstable/http" import { layerWorkersAi } from "./WorkersAiHttpClient" @@ -97,9 +97,53 @@ export interface LlmEnv extends Record { readonly MAPLE_LLM_PROVIDER?: string readonly MAPLE_TRIAGE_MODEL_OPENROUTER?: string readonly MAPLE_TRIAGE_MODEL_WORKERS_AI?: string + /** Context window in tokens, overriding {@link MODEL_LIMITS} for the configured model. */ + readonly MAPLE_TRIAGE_MODEL_CONTEXT?: string + /** Max completion tokens, overriding {@link MODEL_LIMITS} for the configured model. */ + readonly MAPLE_TRIAGE_MODEL_OUTPUT?: string readonly OPENROUTER_API_KEY?: string } +/** + * Context windows for the models Maple configures. + * + * `@maple/llm` has a `ModelLimits { context, output }` on `Model.defaults`, but **no provider + * populates it** — it is `undefined` for every model the vendored package builds. Maple needs it to + * know when a transcript is approaching the wall, so it is filled in here, at the Maple seam, via + * `Model.update`. Nothing is added to `lib/llm` (see `lib/llm/MAPLE.md`). + * + * Conservative on purpose. A limit set too low compacts early, which costs a summarization call; a + * limit set too high overflows, which costs the whole turn. When in doubt, go low. + */ +const MODEL_LIMITS: Record = { + // Verified against OpenRouter's public model catalogue. + "openai/gpt-5.6-luna": { context: 1_050_000, output: 128_000 }, + // Moonshot's own `kimi-k2.6` is 262_144, but Cloudflare does not publish the window its Workers + // AI deployment actually serves, and a serving deployment is usually narrower than upstream's + // maximum. Held at the conservative default until someone measures it; override with + // `MAPLE_TRIAGE_MODEL_CONTEXT` if the real figure turns out to be higher. + "@cf/moonshotai/kimi-k2.6": { context: 128_000, output: 8_000 }, +} + +/** For a model not in the table at all. Low enough that an unknown model compacts rather than fails. */ +const DEFAULT_MODEL_LIMITS = { context: 128_000, output: 8_000 } as const + +const readPositiveInt = (env: LlmEnv, key: keyof LlmEnv): number | undefined => { + const raw = readString(env, key) + if (raw === undefined) return undefined + const value = Number(raw) + return Number.isSafeInteger(value) && value > 0 ? value : undefined +} + +/** + * The context budget a turn should plan against, in tokens, or `undefined` if the model declares + * none. Callers treat `undefined` as "don't compact" rather than guessing a number. + */ +export const contextLimitOf = (model: Model): number | undefined => model.defaults?.limits?.context + +/** Max completion tokens, used to reserve headroom when deciding whether the input still fits. */ +export const outputLimitOf = (model: Model): number | undefined => model.defaults?.limits?.output + const readString = (env: LlmEnv, key: keyof LlmEnv): string | undefined => { const value = env[key] return typeof value === "string" && value.trim() !== "" ? value.trim() : undefined @@ -126,16 +170,39 @@ export const resolveLlmProvider = (env: LlmEnv): LlmProvider => * OpenRouter's body fields and have no meaning to Cloudflare. */ export const resolveTriageModel = (env: LlmEnv, tags?: LlmCallTags): Model => - resolveLlmProvider(env) === "workers-ai" - ? CloudflareWorkersAI.configure({ - accountId: readString(env, "CLOUDFLARE_ACCOUNT_ID") ?? BINDING_PLACEHOLDER, - apiKey: readString(env, "CLOUDFLARE_API_KEY") ?? BINDING_PLACEHOLDER, - }).model(readString(env, "MAPLE_TRIAGE_MODEL_WORKERS_AI") ?? DEFAULT_WORKERS_AI_MODEL) - : OpenRouter.configure({ - apiKey: readString(env, "OPENROUTER_API_KEY") ?? "", - headers: { "HTTP-Referer": OPENROUTER_APP_URL, "X-Title": OPENROUTER_APP_TITLE }, - ...(tags === undefined ? {} : { http: { body: openRouterTagBody(tags) } }), - }).model(readString(env, "MAPLE_TRIAGE_MODEL_OPENROUTER") ?? DEFAULT_OPENROUTER_MODEL) + withLimits( + env, + resolveLlmProvider(env) === "workers-ai" + ? CloudflareWorkersAI.configure({ + accountId: readString(env, "CLOUDFLARE_ACCOUNT_ID") ?? BINDING_PLACEHOLDER, + apiKey: readString(env, "CLOUDFLARE_API_KEY") ?? BINDING_PLACEHOLDER, + }).model(readString(env, "MAPLE_TRIAGE_MODEL_WORKERS_AI") ?? DEFAULT_WORKERS_AI_MODEL) + : OpenRouter.configure({ + apiKey: readString(env, "OPENROUTER_API_KEY") ?? "", + headers: { "HTTP-Referer": OPENROUTER_APP_URL, "X-Title": OPENROUTER_APP_TITLE }, + ...(tags === undefined ? {} : { http: { body: openRouterTagBody(tags) } }), + }).model(readString(env, "MAPLE_TRIAGE_MODEL_OPENROUTER") ?? DEFAULT_OPENROUTER_MODEL), + ) + +/** + * Attach the model's context window, which the vendored providers leave unset. + * + * Purely additive — `defaults.limits` was `undefined` before — so nothing that ignores it can + * regress. `defaults` is spread rather than replaced so a provider that *does* set generation or + * HTTP defaults keeps them. + */ +const withLimits = (env: LlmEnv, model: Model): Model => { + const known = MODEL_LIMITS[String(model.id)] ?? DEFAULT_MODEL_LIMITS + return Model.update(model, { + defaults: { + ...model.defaults, + limits: { + context: readPositiveInt(env, "MAPLE_TRIAGE_MODEL_CONTEXT") ?? known.context, + output: readPositiveInt(env, "MAPLE_TRIAGE_MODEL_OUTPUT") ?? known.output, + }, + }, + }) +} /** * The runnable LLM stack — identical for both providers, which is what makes the switch a pure diff --git a/apps/api/src/workflows/triage-agent.ts b/apps/api/src/workflows/triage-agent.ts index 5de896088..320c14825 100644 --- a/apps/api/src/workflows/triage-agent.ts +++ b/apps/api/src/workflows/triage-agent.ts @@ -22,12 +22,11 @@ import type { AiTriageIncidentKind } from "@maple/domain/http" import { AiTriageResult } from "@maple/domain/http" import { LLM, LLMEvent, Message, ToolResultPart, type LLMRequest, type Model, type Usage } from "@maple/llm" -import { Tool, ToolFailure, ToolRuntime, toDefinitions, type Tools } from "@maple/llm" +import { ToolRuntime, toDefinitions } from "@maple/llm" import { Effect } from "effect" -import { toLlmCallError } from "@/platform/Llm" -import { callMcpTool } from "@/mcp/dispatcher" -import { CurrentMcpTenant } from "@/mcp/lib/query-warehouse" -import { mapleToolDefinitions, toInputSchema } from "@/mcp/tools/registry" +import { contextLimitOf, outputLimitOf, toLlmCallError } from "@/platform/Llm" +import { buildMapleTools } from "@/mcp/tools/llm-tools" +import { isNearContextLimit, pruneToolResults } from "@/chat/context-budget" import type { TenantContext } from "@/services/auth/tenant-context" import { buildTriageContextMessage, TRIAGE_SYSTEM_PROMPT, TRIAGE_TOOL_NAMES } from "./triage-prompt" @@ -56,69 +55,10 @@ export interface TriageAgentOutput { readonly toolSteps: number } -/** - * Re-pin the service requirements of an MCP tool handler. - * - * `MapleToolDefinition.handler` types its requirements as `any` — a deliberate erasure at the - * boundary between ~57 heterogeneous tool implementations and the single `McpServer.addTool` - * signature (see the comment on the type in `mcp/tools/registry.ts`). `@maple/llm`'s `Tool.make` - * insists on `never`, and `any` is not assignable to `never` in that position, so the erasure has - * to be undone somewhere. Doing it here, once and by name, keeps it visible: the services really - * are supplied, by the `ManagedRuntime` the workflow builds around this loop. - */ -// eslint-disable-next-line @typescript-eslint/no-explicit-any -const withRuntimeServices = (effect: Effect.Effect): Effect.Effect => - effect as Effect.Effect - -/** - * Serialize an MCP tool result for the model. Maple's tools already return model-facing text - * blocks, so this is a join rather than a re-encode; `isError` is surfaced as a `ToolFailure` so - * the runtime emits a `tool-error` event and the model can self-correct. - */ -const toolResultText = (result: { content: ReadonlyArray<{ text: string }> }): string => - result.content.map((block) => block.text).join("\n") - -/** - * Wrap the Maple MCP registry as `@maple/llm` tools, restricted to the read-only allowlist. - * - * Dynamic (`jsonSchema`) mode is the right fit: `toInputSchema` already produces the exact JSON - * Schema the MCP surface publishes, and `callMcpTool` does its own Effect Schema decode with the - * tool's own error messages. Decoding twice would only give the model a second, worse phrasing of - * the same validation failure. - * - * The tenant is provided per call rather than ambiently so the loop can never widen its own scope: - * every tool executes under exactly the org the workflow was enqueued for. - */ -const buildTriageTools = (tenant: TenantContext): Tools => - Object.fromEntries( - mapleToolDefinitions - .filter((definition) => TRIAGE_TOOL_NAMES.has(definition.name)) - .map((definition) => [ - definition.name, - Tool.make({ - description: definition.description, - jsonSchema: toInputSchema(definition.schema), - execute: (params): Effect.Effect => - withRuntimeServices( - callMcpTool(definition.name, params).pipe( - Effect.provideService(CurrentMcpTenant, tenant), - Effect.flatMap((result) => - result.isError - ? Effect.fail(new ToolFailure({ message: toolResultText(result) })) - : Effect.succeed(toolResultText(result)), - ), - // A tool that fails outright (unknown tool, tenant error) must not kill - // the investigation — hand the model the message and let it route around. - Effect.catchCause((cause) => - Effect.fail( - new ToolFailure({ message: `Tool failed: ${String(cause)}` }), - ), - ), - ), - ), - }), - ]), - ) +/** The read-only allowlist, as `@maple/llm` tools. Nothing here is approval-gated: the triage loop + * is autonomous and mutating tools are simply not in `TRIAGE_TOOL_NAMES`. */ +const buildTriageTools = (tenant: TenantContext) => + buildMapleTools(tenant, { include: (name) => TRIAGE_TOOL_NAMES.has(name) }) const addUsage = (total: { input: number; output: number; cacheRead: number }, usage: Usage | undefined) => ({ input: total.input + (usage?.inputTokens ?? 0), @@ -184,6 +124,17 @@ export const runTriageAgent = Effect.fn("ai_triage.investigate")(function* (inpu ), ], }) + + // Same exposure as the chat turn, and worse: twelve steps of warehouse-sized tool payloads + // with no user in the loop to notice it stalling. Acts on the provider's reported count. + if ( + isNearContextLimit(response.usage?.inputTokens ?? 0, { + context: contextLimitOf(input.model), + output: outputLimitOf(input.model), + }) + ) { + request = pruneToolResults(request) + } } yield* Effect.annotateCurrentSpan({ diff --git a/apps/web/src/components/ai-elements/types.ts b/apps/web/src/components/ai-elements/types.ts index eb84ef0b4..f86a708f3 100644 --- a/apps/web/src/components/ai-elements/types.ts +++ b/apps/web/src/components/ai-elements/types.ts @@ -52,6 +52,24 @@ export type UIMessagePart = input: unknown errorText: string } + /** + * A sub-agent run: its own transcript, nested under the `task` call that started it. + * + * A part of its own rather than a `dynamic-tool` with a payload, because it renders as a + * collapsible transcript rather than a tool row, and because `transcript-rows` must not fold it + * into a "Used N tools" group — a sub-agent is content, not plumbing. + * + * `messages` is `UIMessage[]` and not recursive by accident: sub-agents cannot spawn sub-agents, + * so a nested message never carries another `task` part. + */ + | { + type: "task" + toolCallId: string + agent: string + description: string + status: "running" | "completed" | "error" | "aborted" + messages: UIMessage[] + } export interface UIMessage { id: string diff --git a/apps/web/src/components/chat/chat-transcript.test.tsx b/apps/web/src/components/chat/chat-transcript.test.tsx index a491c6860..1d86220d5 100644 --- a/apps/web/src/components/chat/chat-transcript.test.tsx +++ b/apps/web/src/components/chat/chat-transcript.test.tsx @@ -1,6 +1,6 @@ // @vitest-environment jsdom -import { cleanup, render, screen } from "@testing-library/react" +import { cleanup, fireEvent, render, screen } from "@testing-library/react" import { afterEach, describe, expect, it, vi } from "vitest" import { ChatTranscript, findDiagnosisMessageId } from "./chat-transcript" @@ -244,3 +244,52 @@ describe("machine-written turns", () => { expect(screen.queryByText(/subject: api/)).toBeNull() }) }) + +describe("ChatTranscript sub-agent cards", () => { + const taskMessage = (status: "running" | "completed" = "completed"): UIMessage => + ({ + id: "m1", + role: "assistant", + parts: [ + { + type: "task", + toolCallId: "t1", + agent: "explore", + description: "trace checkout latency", + status, + messages: [ + { + id: "c1", + role: "assistant", + parts: [{ type: "text", text: "p99 is 4.2s in checkout-api.", state: "done" }], + }, + ], + }, + ], + }) as unknown as UIMessage + + it("renders a collapsed card naming the sub-agent and what it was asked", () => { + render() + + expect(screen.getByText("explore")).toBeTruthy() + expect(screen.getByText("trace checkout latency")).toBeTruthy() + // Collapsed by default: the point of delegating is that the parent thread does not carry + // the sub-agent's search. + expect(screen.queryByText("p99 is 4.2s in checkout-api.")).toBeNull() + }) + + it("expands to the sub-agent's own transcript on click", () => { + render() + + fireEvent.click(screen.getByText("explore")) + expect(screen.getByText("p99 is 4.2s in checkout-api.")).toBeTruthy() + }) + + it("never folds a sub-agent into a Used N tools header", () => { + // A sub-agent run is content, not plumbing. + render() + + expect(screen.queryByText(/Used \d+ tools/)).toBeNull() + expect(items().map((el) => (el as HTMLElement).dataset.messageId)).toEqual(["m1"]) + }) +}) diff --git a/apps/web/src/components/chat/chat-transcript.tsx b/apps/web/src/components/chat/chat-transcript.tsx index 0d5bf177d..b485598e5 100644 --- a/apps/web/src/components/chat/chat-transcript.tsx +++ b/apps/web/src/components/chat/chat-transcript.tsx @@ -21,6 +21,7 @@ import { StatusMarker } from "@/components/ai-elements/status-marker" import { Tool, ToolRow, toolLabel } from "@/components/ai-elements/tool" import { ToolGroup } from "@/components/ai-elements/tool-group" import { ApprovalCard } from "./approval-card" +import { TaskCard } from "./task-card" import { DiagnosisReportCard } from "./diagnosis-report-card" import { MessageActions, messageText } from "./message-actions" import { parseDiagnosisMarker } from "./diagnosis-marker" @@ -148,6 +149,20 @@ function renderMessageParts({ if (text) nodes.push({text}) continue } + // A sub-agent run is content, not plumbing: its own card, never folded into a tool group. + if (part.type === "task") { + flushTools() + nodes.push( + , + ) + continue + } if (!isToolPart(part)) continue const tp = part as ToolPart diff --git a/apps/web/src/components/chat/task-card.tsx b/apps/web/src/components/chat/task-card.tsx new file mode 100644 index 000000000..28649f2e6 --- /dev/null +++ b/apps/web/src/components/chat/task-card.tsx @@ -0,0 +1,108 @@ +import { useState } from "react" +import { + ChatBubbleSparkleIcon, + ChevronDownIcon, + ChevronRightIcon, + CircleCheckIcon, + CircleXmarkIcon, + LoaderIcon, +} from "@/components/icons" +import { Tool } from "@/components/ai-elements/tool" +import { RichText } from "@/components/ai-elements/rich-text" +import type { UIMessage } from "@/components/ai-elements/types" +import { toolNameFor, type ToolPart } from "./transcript-rows" + +type TaskStatus = "running" | "completed" | "error" | "aborted" + +interface TaskCardProps { + agent: string + description: string + status: TaskStatus + messages: readonly UIMessage[] +} + +const STATUS_ICON: Record = { + running: LoaderIcon, + completed: CircleCheckIcon, + error: CircleXmarkIcon, + aborted: CircleXmarkIcon, +} + +const STATUS_TINT: Record = { + running: "text-muted-foreground", + completed: "text-success", + error: "text-destructive", + aborted: "text-muted-foreground", +} + +const toolCount = (messages: readonly UIMessage[]): number => + messages.reduce( + (total, message) => total + message.parts.filter((part) => part.type === "dynamic-tool").length, + 0, + ) + +/** + * One sub-agent run, collapsed. + * + * Collapsed by default because the whole point of delegating is that the parent conversation does + * not have to carry the sub-agent's search. Its answer already reached the model through the tool + * result; this card is for a reader who wants to check the work. + */ +export function TaskCard({ agent, description, status, messages }: TaskCardProps) { + const [open, setOpen] = useState(false) + const Icon = STATUS_ICON[status] + const tools = toolCount(messages) + + return ( +
+ + + {open ? ( +
+ {messages.length === 0 ? ( +

Nothing yet.

+ ) : ( + messages.map((message) => ( +
+ {message.parts.map((part, index) => + part.type === "text" ? ( + {part.text} + ) : part.type === "dynamic-tool" ? ( + + ) : null, + )} +
+ )) + )} +
+ ) : null} +
+ ) +} diff --git a/apps/web/src/components/chat/transcript-rows.test.ts b/apps/web/src/components/chat/transcript-rows.test.ts index 0b31da21c..c43d3fda6 100644 --- a/apps/web/src/components/chat/transcript-rows.test.ts +++ b/apps/web/src/components/chat/transcript-rows.test.ts @@ -93,3 +93,31 @@ describe("buildTranscriptRows", () => { expect(isToolOnlyMessage(text("u", "user", "hi"))).toBe(false) }) }) + +describe("sub-agent parts", () => { + const task = (id: string): UIMessage => + ({ + id, + role: "assistant", + parts: [ + { + type: "task", + toolCallId: `${id}-t`, + agent: "explore", + description: "trace checkout latency", + status: "completed", + messages: [], + }, + ], + }) as unknown as UIMessage + + it("is not tool-only, so it never disappears into a Used N tools header", () => { + // A sub-agent run is content: the reader delegated part of the investigation and should see + // that it happened, not have it folded away as plumbing. + expect(isToolOnlyMessage(task("m1"))).toBe(false) + }) + + it("breaks a run of tool-only turns", () => { + expect(kinds([tools("m1", 2), task("m2"), tools("m3", 2)])).toEqual(["message", "message", "message"]) + }) +}) diff --git a/apps/web/src/hooks/use-maple-chat.test.tsx b/apps/web/src/hooks/use-maple-chat.test.tsx index 95d1f8e09..13e59d1ff 100644 --- a/apps/web/src/hooks/use-maple-chat.test.tsx +++ b/apps/web/src/hooks/use-maple-chat.test.tsx @@ -57,7 +57,13 @@ function Harness() { {chat.messages .flatMap((message) => message.parts) - .map((part) => (part.type === "text" ? part.text : `[${part.state}]`)) + .map((part) => + part.type === "text" + ? part.text + : part.type === "task" + ? `[task:${part.agent}:${part.status}]` + : `[${part.state}]`, + ) .join("|")}