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 deleted file mode 100644 index 48e002126..000000000 --- a/apps/api/src/chat/agent.test.ts +++ /dev/null @@ -1,178 +0,0 @@ -/** - * `runChatTurn` — the event stream one submission produces. - * - * The model is a stub layer over `LLMClient`, so these assert the *turn loop*: what it emits, in - * what order, and — the part that shipped wrong — what it does NOT emit once the turn is over. - * Every failure mode here was invisible to `tsc` and to the branch's suite. - */ -import { assert, describe, it } from "vitest" -import { Effect, Layer, Stream } from "effect" -import { LLM, LLMClient, LLMEvent, type LLMRequest, type Model } from "@maple/llm" -import { CloudflareWorkersAI } from "@maple/llm/providers/cloudflare" -import { runChatTurn, type ChatTurnEvent } from "./agent" -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: "test", - apiKey: "test", -}).model("@cf/test/model") - -/** A text delta, as the provider-neutral event the turn folds. */ -const textDelta = (text: string): LLMEvent => ({ type: "text-delta", id: "t1", text }) as LLMEvent - -const finish = (): LLMEvent => ({ type: "finish", reason: "stop" }) as LLMEvent - -const toolCall = (id: string, name: string): LLMEvent => - ({ type: "tool-call", id, name, input: {}, providerExecuted: false }) as LLMEvent - -/** - * Stub the model with a scripted event stream per step. - * - * `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 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 } - -const stubModel = (steps: ReadonlyArray) => { - 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) => { - const scripted = steps[step] ?? [finish()] - step += 1 - // The vendored error shape the turn maps through `toLlmCallError`. - const failure = { - _tag: "LLMError", - module: "test", - method: "stream", - reason: { _tag: "ProviderInternal" }, - message: "upstream exploded", - retryable: true, - } 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) -} - -const collect = ( - steps: ReadonlyArray, - overrides: { readonly sessionId?: string; readonly isCurrent?: () => boolean } = {}, -) => - runChatTurn({ - sessionId: overrides.sessionId ?? "org_test:tab", - tenant: TENANT, - model: MODEL, - messages: [], - messageId: "m1", - ...(overrides.isCurrent ? { isCurrent: overrides.isCurrent } : {}), - }).pipe( - Stream.runCollect, - Effect.map((events) => Array.from(events) as ChatTurnEvent[]), - Effect.provide(stubModel(steps)), - ) - -const types = (events: ReadonlyArray) => events.map((event) => event.type) - -const terminal = (events: ReadonlyArray) => events.filter((event) => event.type === "turn-end") - -describe("runChatTurn", () => { - it("emits turn-start, the text, and exactly one turn-end", async () => { - const events = await Effect.runPromise(collect([[textDelta("Hello"), finish()]])) - - assert.deepEqual(types(events), ["turn-start", "text-delta", "turn-end"]) - assert.lengthOf(terminal(events), 1) - }) - - 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()]])) - - // 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. - const deltas = events.filter((event) => event.type === "text-delta") - assert.lengthOf(deltas, 1) - assert.equal( - deltas.map((event) => (event.type === "text-delta" ? event.text : "")).join(""), - chunks.join(""), - ) - }) - - it("keeps text ahead of the tool calls it precedes", async () => { - const events = await Effect.runPromise( - collect([ - [textDelta("Looking"), textDelta(" it up"), toolCall("c1", "create_alert_rule"), finish()], - ]), - ) - - // Batching must never reorder: the deltas live in one stream segment and the tool events in - // the concatenated one after it, so a slow batch cannot overtake the call it introduced. - assert.deepEqual(types(events), ["turn-start", "text-delta", "tool-call", "turn-end"]) - const delta = events.find((event) => event.type === "text-delta") - 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 }])) - - // 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. - // Both landed in the durable log; the SSE route stops at the first, so it only surfaced on - // the next reload. - assert.lengthOf(terminal(events), 1) - const end = terminal(events)[0] - assert.equal(end?.type === "turn-end" ? end.reason : undefined, "error") - }) - - it("dispatches NO tools when the stream fails after announcing one", async () => { - // 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 }]), - ) - - assert.isEmpty( - events.filter((event) => event.type === "tool-result"), - "a failed turn must not run tools past its terminal event", - ) - }) - - 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 proposals = events.filter((event) => event.type === "tool-call") - assert.lengthOf(proposals, 1) - assert.equal( - proposals[0]?.type === "tool-call" ? proposals[0].proposed : undefined, - true, - "a gated call is a proposal, not an execution", - ) - // Nothing fabricates an outcome: the model is never told the mutation happened. - assert.isEmpty(events.filter((event) => event.type === "tool-result")) - assert.lengthOf(terminal(events), 1) - }) - - it("stops without a second terminal event when the turn is superseded", async () => { - const events = await Effect.runPromise( - collect([[textDelta("hi"), finish()]], { isCurrent: () => false }), - ) - - // An abort already recorded the terminal event on the session; emitting another here would - // double-close the turn in the durable log. - assert.isEmpty(terminal(events)) - }) -}) diff --git a/apps/api/src/chat/agent.ts b/apps/api/src/chat/agent.ts deleted file mode 100644 index 08ed93144..000000000 --- a/apps/api/src/chat/agent.ts +++ /dev/null @@ -1,478 +0,0 @@ -/** - * The Maple chat agent turn, in process on `@maple/llm`. - * - * This is the streaming sibling of `workflows/triage-agent.ts`: same tool wrapping, same - * multi-turn loop, but it emits `ChatEvent`s as they happen instead of returning a structured - * result, and it gates mutating tools. - * - * **Approvals are a real interrupt now.** Flue's event stream had no human-in-the-loop primitive, - * so `apps/chat-flue/src/lib/approval.ts` swapped every mutating tool for one whose `execute` - * returned a `{ status: "proposed" }` marker without mutating — a propose-then-apply stub the web - * client then re-ran through `POST /api/chat/apply`. Here the loop simply *stops* on a gated call: - * it emits a `tool-call` event with `proposed: true` and ends the turn. Nothing fabricates a tool - * result, so the model is never told a mutation happened when it did not. - * - * `POST /api/chat/apply` still exists and is still how the approved mutation runs — it is the - * user's action, authenticated as the user, which is exactly where it belongs. - */ -import { - chatModeFromSessionId, - investigationIdFromChatSessionId, - type ChatEvent, -} from "@maple/domain/chat-session" -import { AiTriageResult, SubmitDiagnosisRequest } from "@maple/domain/http" -import { InvestigationId } from "@maple/domain/primitives" -import { - LLM, - LLMEvent, - LLMResponse, - Message, - ToolResultPart, - type LLMRequest, - type Model, - 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 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" - -const decodeInvestigationIdOption = Schema.decodeUnknownOption(InvestigationId) - -/** Hard cap on assistant turns per submission. */ -const MAX_STEPS = 10 - -/** Fan-out cap for tool calls issued in the same assistant turn. */ -const TOOL_CONCURRENCY = 4 - -/** - * How many text deltas are folded into one emitted event, and how long a partial batch waits. - * - * Roughly one animation frame. Every delta that leaves this stream becomes a durable SQLite row, an - * SSE frame and a React state commit, and the browser cannot show more than one update per frame - * anyway — so batching to that granularity costs no perceptible smoothness and removes most of the - * per-token work at all three layers. The size cap keeps a fast provider from letting a batch grow - * unboundedly within the window. - */ -const DELTA_BATCH_SIZE = 24 -const DELTA_BATCH_WINDOW = "16 millis" - -/** - * Running token total for one turn, accumulated across its steps. - * - * Mutable and shared rather than returned, because the one consumer — `submit_diagnosis` — is a - * *tool* invoked mid-turn, so there is no "after the turn" moment at which to hand it a total. - * In practice the diagnosis call is the last thing an investigation does, so this is the whole - * turn bar the final assistant message. Before this, `SubmitDiagnosisRequest` was built with no - * usage at all, so `InvestigationService`'s `if (env && (inputTokens || outputTokens))` was always - * false: `investigations.model` stayed null and Autumn was never metered for autonomous - * investigations, which the pre-`@maple/llm` workflow path did meter. - */ -export interface TurnUsage { - input: number - output: number - cacheRead: number -} - -export const makeTurnUsage = (): TurnUsage => ({ input: 0, output: 0, cacheRead: 0 }) - -const addUsage = (total: TurnUsage, usage: Usage | undefined): void => { - total.input += usage?.inputTokens ?? 0 - total.output += usage?.outputTokens ?? 0 - 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 - readonly model: Model - /** The full transcript so far, oldest first, already including the new user message. */ - readonly messages: ReadonlyArray - readonly messageId: string - /** - * Investigate-mode sessions get a `submit_diagnosis` tool. It is supplied rather than built - * here because it needs `InvestigationService`, which would otherwise drag the service graph - * into this module's imports. - */ - readonly extraTools?: Tools - /** - * Whether this turn still holds the session's turn slot. - * - * Checked between steps so an abort takes effect at the next boundary instead of only after the - * in-flight model call drains, and so a turn that has been superseded stops writing into a - * conversation that has moved on. Defaults to "always current" for callers with no session. - */ - readonly isCurrent?: () => boolean - /** Accumulates this turn's token usage; see {@link TurnUsage}. */ - readonly usage?: TurnUsage -} - -/** - * The `submit_diagnosis` tool for an investigate-mode session (`":inv-"`). - * - * The agent calls it exactly once at the end of its autonomous pass and its arguments ARE the - * structured report — `AiTriageResult` directly, not the Valibot mirror `apps/chat-flue` had to - * keep in sync by hand. - * - * Deliberately not approval-gated: it is the structured-output channel, not a user-facing - * mutation. The investigation id and org ride from the session id, so the agent never chooses - * which investigation it writes. - * - * `submitDiagnosis` arrives as a callback rather than being resolved from `InvestigationService` - * here: that service is itself what starts an investigation's autonomous turn, so resolving it - * through the Effect requirements channel would make `InvestigationService` require itself. - */ -export type SubmitDiagnosis = ( - orgId: TenantContext["orgId"], - investigationId: InvestigationId, - request: SubmitDiagnosisRequest, - // eslint-disable-next-line @typescript-eslint/no-explicit-any -) => Effect.Effect - -export const buildSubmitDiagnosisTool = ( - sessionId: string, - tenant: TenantContext, - submitDiagnosis: SubmitDiagnosis, - usage: TurnUsage, - model: Model, -): Tools => { - const tools: Tools = {} - const rawId = investigationIdFromChatSessionId(sessionId) - if (!rawId) return tools - // `decodeUnknownSync` here turned a session id whose `inv-` suffix was not a UUID into a thrown - // defect on a user-supplied string. An unparseable id simply means this conversation is not an - // investigation, so it gets no `submit_diagnosis` tool. - const decoded = decodeInvestigationIdOption(rawId) - if (Option.isNone(decoded)) return tools - const investigationId = decoded.value - tools.submit_diagnosis = Tool.make({ - description: - "Record your structured diagnosis for THIS investigation. Call it exactly once, " + - "after you have gathered evidence, with your final assessment. It persists the report " + - "and renders it for the user. After calling it, stop unless the user asks a follow-up.", - parameters: AiTriageResult, - success: Schema.String, - execute: (report) => - withRuntimeServices( - submitDiagnosis( - tenant.orgId, - investigationId, - new SubmitDiagnosisRequest({ - report, - model: String(model.id), - inputTokens: usage.input, - outputTokens: usage.output, - }), - ).pipe( - Effect.as("Diagnosis recorded."), - // Named failures only. `catchCause` + `String(cause)` fed the model a rendered - // Effect cause — stack frames, and connection details out of a DatabaseError. - Effect.catchCause((cause) => - Effect.fail( - new ToolFailure({ - message: `submit_diagnosis failed: ${summarizeToolFailure(cause)}`, - }), - ), - ), - ), - ), - }) - 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)}`, - }), - ), - ), - ), - ), - }), - ] - }), - ) - -/** Distributive `Omit`, so each union member keeps its own shape. */ -type WithoutSeq = T extends unknown ? Omit : never - -/** - * Events this turn wants appended to the session log. `seq` is assigned by the Durable Object, - * which owns the ordering, so the agent emits everything without one. `user-message` is excluded: - * the session records the user turn at submission time, before the agent runs. - */ -export type ChatTurnEvent = WithoutSeq> - -/** - * Run one submission to completion, streaming `ChatTurnEvent`s. - * - * The model is *streamed*, not `generate`d, so text deltas reach the session log — and through it - * the client — while the turn is still running. The raw `LLMEvent`s are folded into an - * `LLMResponse` on the way past so the assistant turn can be appended to the transcript verbatim - * for the next step. - */ -export const runChatTurn = (input: ChatTurnInput): Stream.Stream => - Stream.unwrap( - Effect.sync(() => { - const tools = { ...buildChatTools(input.tenant), ...input.extraTools } - const system = buildSystemPrompt({ mode: chatModeFromSessionId(input.sessionId) }) - const request = LLM.request({ - id: input.messageId, - model: input.model, - system, - 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 turnEnd = ( - input: ChatTurnInput, - reason: Extract["reason"], - error?: string, -): ChatTurnEvent => ({ - 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() - -/** - * One assistant turn, then either settle its tool calls and recurse, or stop. - * - * Recursion rather than a loop because each step's output is a `Stream` that must be concatenated - * lazily: the next request cannot be built until the current step's tool results exist. - */ -const runStep = ( - input: ChatTurnInput, - tools: Tools, - request: LLMRequest, - step: number, -): Stream.Stream => - Stream.suspend(() => { - 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 - // *second* terminal event after the error one — and, if that partial response carried tool - // calls, would dispatch them and recurse after the turn had already been declared over. - // Those extra events land invisibly (the SSE route stops at the first `turn-end`) and - // surface on the next reload. - let failed = false - - const live: Stream.Stream = LLM.stream(request).pipe( - Stream.tap((event) => Effect.sync(() => collected.push(event))), - Stream.filter((event) => event.type === "text-delta" && event.text !== ""), - // One durable row, one SSE frame and one React commit per *token* is more fidelity than - // a screen can show. Coalescing into roughly one frame's worth of deltas is invisible - // to a reader and cuts all three by about an order of magnitude. Only text deltas are - // batched, and only against each other — `collected` still holds every raw event, and - // 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.catch((error) => { - failed = true - return Stream.fromIterable([ - turnEnd(input, "error", toLlmCallError("chat.turn", error).message), - ]) - }), - ) - - const settleAndRecurse = Stream.unwrap( - Effect.gen(function* () { - if (failed) return Stream.empty - // Aborted between steps: the session already recorded the terminal event, so stop - // without emitting a second one. - if (!isCurrent(input)) return Stream.empty - - const response = LLMResponse.fromEvents(collected) - // A stream that neither failed nor assembled still ended the turn; say so, rather - // than leaving the log with no terminal event at all. - if (!response) return Stream.fromIterable([turnEnd(input, "stop")]) - - if (input.usage) addUsage(input.usage, response.usage) - - const calls = response.events - .filter(LLMEvent.is.toolCall) - .filter((call) => !call.providerExecuted) - - if (calls.length === 0) return Stream.fromIterable([turnEnd(input, "stop")]) - - // 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)) - if (gated) { - const proposal: ChatTurnEvent = { - type: "tool-call", - 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", - messageId: input.messageId, - callId: call.id, - name: call.name, - input: call.input, - }), - ) - - // Announce first, settle second, as two stream segments. Emitting both together - // after `Effect.forEach` resolved meant a tool call only ever reached the log - // *already finished*, so the UI could never render one running — most of the point - // of streaming a turn that spends its time in tools. - const settled = Stream.unwrap( - Effect.gen(function* () { - const dispatched = yield* Effect.forEach( - calls, - (call) => - ToolRuntime.dispatch(tools, call).pipe( - Effect.map((result) => [call, result] as const), - ), - { concurrency: TOOL_CONCURRENCY }, - ) - - const results = dispatched.map( - ([call, outcome]): ChatTurnEvent => ({ - type: "tool-result", - messageId: input.messageId, - callId: call.id, - output: outcome.result.value, - ...(outcome.result.type === "error" ? { isError: true } : {}), - }), - ) - - // Aborted while the tools were in flight: record what they returned so the - // 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 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, - }), - ), - ), - ], - }) - return Stream.concat( - Stream.fromIterable(results), - runStep(input, tools, next, step + 1), - ) - }), - ) - - return Stream.concat(Stream.fromIterable(announced), settled) - }), - ) - - return Stream.concat(live, settleAndRecurse) - }) 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/loop/budgets.ts b/apps/api/src/chat/loop/budgets.ts new file mode 100644 index 000000000..effe0d360 --- /dev/null +++ b/apps/api/src/chat/loop/budgets.ts @@ -0,0 +1,125 @@ +/** + * Every ceiling the turn loop enforces, in one place. + * + * Collected deliberately. These numbers only make sense against each other — the per-step attempt + * cap multiplies against the step cap, which multiplies against the sub-agent fan-out, and the + * product has to stay under `ChatSession`'s `TURN_STALE_MS` watchdog or the Durable Object will + * declare a *still-running* turn abandoned and write a terminal event underneath it. Scattered + * across three modules, that relationship was invisible and each constant looked independently + * reasonable. + * + * The worst case, spelled out: `MAX_STEPS` (10) × `TASK_BUDGET_PER_TURN` (4) × `SUBAGENT_MAX_STEPS` + * (6) = 240 model calls, which `TURN_STEP_BUDGET` (30) is the backstop against. + */ +import { Semaphore } from "effect" + +// --------------------------------------------------------------------------- +// Steps +// --------------------------------------------------------------------------- + +/** + * 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, + * rather than stopping dead on a wall of tool rows with no words. + */ +export const MAX_STEPS = 10 + +/** Fan-out cap for tool calls issued in the same assistant turn. */ +export const TOOL_CONCURRENCY = 4 + +/** + * 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 + +/** + * Model calls across the parent turn and every descendant. + * + * The backstop against the multiplication above. If real turns start approaching this, lower + * `TASK_BUDGET_PER_TURN` — do **not** raise `TURN_STALE_MS` to compensate. + */ +export const TURN_STEP_BUDGET = 30 + +// --------------------------------------------------------------------------- +// Retry +// --------------------------------------------------------------------------- + +/** 1 initial attempt + 3 retries. */ +export const MAX_STEP_ATTEMPTS = 4 + +export const STEP_RETRY_BASE_MS = 1_000 +export 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. + */ +export 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 the session's single turn slot the entire time. + */ +export const STEP_RETRY_BUDGET_MS = 60_000 + +// --------------------------------------------------------------------------- +// Delegation +// --------------------------------------------------------------------------- + +/** + * How deep nesting may go. 1 means the conversation's turn may spawn, and sub-agents may not. + * + * Capped twice, as opencode does: structurally, because every sub-agent's ruleset denies `task` so + * the tool is never offered; and numerically, here. Belt and braces, because the structural cap + * silently disappears the moment someone writes a sub-agent ruleset with `{"*": allow}`. + */ +export const SUBAGENT_MAX_DEPTH = 1 + +/** Concurrent sub-agent turns inside one Durable Object. */ +export const TASK_CONCURRENCY = 2 + +/** Total sub-agent turns one parent turn may start, across all its steps. */ +export const TASK_BUDGET_PER_TURN = 4 + +// --------------------------------------------------------------------------- +// The mutable records +// --------------------------------------------------------------------------- + +/** + * Time already spent in backoff this turn. + * + * Mutable and shared rather than returned, the same shape and for the same reason as `TurnUsage`: + * the consumer is mid-turn, so there is no "after the turn" moment to reconcile at. A per-step cap + * would compose badly — ten steps each retrying three times is minutes of pure backoff. + */ +export interface StepRetryBudget { + spentMs: number +} + +export const makeStepRetryBudget = (): StepRetryBudget => ({ spentMs: 0 }) + +/** Delegation state, shared by the parent turn and every descendant. */ +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 make another model call at all. Checked by the loop and by `delegate`. */ +export const hasStepBudget = (budget: TaskBudget | undefined): boolean => + budget === undefined || budget.stepsUsed < TURN_STEP_BUDGET diff --git a/apps/api/src/chat/loop/context.test.ts b/apps/api/src/chat/loop/context.test.ts new file mode 100644 index 000000000..a5b8bbb9b --- /dev/null +++ b/apps/api/src/chat/loop/context.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" + +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/loop/context.ts b/apps/api/src/chat/loop/context.ts new file mode 100644 index 000000000..96ef2cc89 --- /dev/null +++ b/apps/api/src/chat/loop/context.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/loop/delegate.test.ts b/apps/api/src/chat/loop/delegate.test.ts new file mode 100644 index 000000000..2fcd97b72 --- /dev/null +++ b/apps/api/src/chat/loop/delegate.test.ts @@ -0,0 +1,344 @@ +/** + * 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 "./index" +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) { + // `compaction` is turn-runner bookkeeping and never reaches this sink; every other + // member carries the ref. + assert.notEqual(event.type, "compaction") + assert.equal( + event.type === "compaction" ? 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") + }) +}) + +describe("the task tool's wire schema", () => { + it("compiles to JSON Schema a provider will accept", async () => { + // The closest thing to provider-contract verification available without a live call. Every + // other test here stubs the model, so nothing else would notice if `Schema.Literals` over a + // runtime-built name list produced something a provider rejects. + let schema: Record | undefined + const service = { + prepare: () => Effect.die(new Error("unused")), + generate: () => Effect.die(new Error("unused")), + stream: (request: LLMRequest) => { + schema = request.tools.find((tool) => tool.name === "task")?.inputSchema as + | Record + | undefined + 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))), + ) + + assert.equal(schema?.type, "object") + assert.deepEqual(schema?.required, ["description", "prompt", "subagent_type"]) + assert.equal(schema?.additionalProperties, false) + + // `subagent_type` must constrain the model to real agent names — a free-form string here + // would let it invent one and turn every delegation into a tool failure. + const properties = schema?.properties as Record> + const subagentType = properties.subagent_type! + const enumValues = + (subagentType.enum as ReadonlyArray | undefined) ?? + (subagentType.anyOf as ReadonlyArray<{ enum?: ReadonlyArray }> | undefined)?.[0]?.enum ?? + [] + assert.include([...enumValues], "explore") + + // Known wart, pinned rather than fixed: Effect renders the literal union as a single-member + // `anyOf` wrapping the enum instead of a bare enum. Harmless on OpenRouter and Workers AI, + // but it is exactly the shape opencode's `tool/json-schema.ts` normalizes away because some + // providers reject a degenerate `anyOf`. If a provider ever does, collapse it here. + assert.isDefined(subagentType.anyOf ?? subagentType.enum) + }) +}) diff --git a/apps/api/src/chat/loop/delegate.ts b/apps/api/src/chat/loop/delegate.ts new file mode 100644 index 000000000..b72507ba7 --- /dev/null +++ b/apps/api/src/chat/loop/delegate.ts @@ -0,0 +1,191 @@ +/** + * `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 caps this respects — depth, fan-out, per-turn count — live in `./budgets.ts` alongside every + * other ceiling, because they only make sense multiplied against the step caps. + * + * ## 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, Stream, Effect } from "effect" +import { Tool, ToolFailure, type Tools } from "@maple/llm" +import { Message } from "@maple/llm" +import { AGENTS, type AgentDefinition } from "../agents" +import { hasStepBudget, SUBAGENT_MAX_DEPTH, TASK_BUDGET_PER_TURN, type TaskBudget } from "./budgets" +import type { ChatTurnEvent, ChatTurnInput } from "./types" + +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 `./turn.ts` while + * `./turn.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/loop/index.ts b/apps/api/src/chat/loop/index.ts new file mode 100644 index 000000000..6f68e177b --- /dev/null +++ b/apps/api/src/chat/loop/index.ts @@ -0,0 +1,25 @@ +/** + * The chat turn loop. + * + * Split out of the surrounding `chat/` directory so the control flow can be read on its own. The + * rest of `chat/` is the world the loop runs in — the Durable Object that owns the log and the turn + * slot, the agent registry, the permission rulesets, the prompt text, the tool set. None of that + * decides how a turn *progresses*; all of it is input. + * + * Read in this order: + * + * 1. `turn.ts` — the control flow. Start here; it is the only file that decides anything. + * 2. `types.ts` — the vocabulary: what goes in, what comes out, what one step carries. + * 3. `budgets.ts` — every ceiling, together, because they multiply against each other. + * 4. `retry.ts` — which failures earn another attempt. + * 5. `context.ts` — keeping the request inside the model's window. + * 6. `delegate.ts` — handing a sub-question to a sub-agent, which re-enters `turn.ts`. + * + * Callers should import from here rather than reaching into a file, so the internal split stays + * free to move. + */ +export { runChatTurn } from "./turn" +export { makeTurnUsage, type ChatTurnEvent, type ChatTurnInput, type TurnUsage } from "./types" +export { MAX_STEPS, SUBAGENT_MAX_STEPS } from "./budgets" +export { isNearContextLimit, pruneToolResults } from "./context" +export { isRetryableStepFailure, stepRetryDelayMs } from "./retry" diff --git a/apps/api/src/chat/loop/retry.test.ts b/apps/api/src/chat/loop/retry.test.ts new file mode 100644 index 000000000..5ec3518f4 --- /dev/null +++ b/apps/api/src/chat/loop/retry.test.ts @@ -0,0 +1,70 @@ +/** + * 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 `./retry.ts`. + */ +import { LlmCallError } from "@maple/domain/llm" +import { assert, describe, it } from "vitest" +import { MAX_STEP_ATTEMPTS } from "./budgets" +import { isRetryableStepFailure, stepRetryDelayMs } from "./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/loop/retry.ts b/apps/api/src/chat/loop/retry.ts new file mode 100644 index 000000000..768f8a732 --- /dev/null +++ b/apps/api/src/chat/loop/retry.ts @@ -0,0 +1,50 @@ +/** + * Which model-call failures a chat step retries, and how long it waits. + * + * Pure policy, deliberately separate from the loop that applies it. `./turn.ts` retries a + * *stream*, so it needs the retraction machinery around the emitted deltas, while + * `apps/api/src/workflows/triage-agent.ts` retries an `Effect` and could use `Effect.retry` + * directly. The classification must not diverge between them, so it lives here. + + * The ceilings this policy is applied under live in `./budgets.ts`, with every other limit the + * loop enforces. + * + * ## 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" + +import { STEP_RETRY_BASE_MS, STEP_RETRY_FACTOR, STEP_RETRY_MAX_MS } from "./budgets" + +/** + * 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/loop/turn.test.ts b/apps/api/src/chat/loop/turn.test.ts new file mode 100644 index 000000000..1e27709e3 --- /dev/null +++ b/apps/api/src/chat/loop/turn.test.ts @@ -0,0 +1,470 @@ +/** + * `runChatTurn` — the event stream one submission produces. + * + * The model is a stub layer over `LLMClient`, so these assert the *turn loop*: what it emits, in + * what order, and — the part that shipped wrong — what it does NOT emit once the turn is over. + * Every failure mode here was invisible to `tsc` and to the branch's suite. + */ +import { assert, describe, it } from "vitest" +import { Effect, Layer, Stream } from "effect" +import { LLMClient, LLMEvent, type LLMRequest, type Model } from "@maple/llm" +import { CloudflareWorkersAI } from "@maple/llm/providers/cloudflare" +import { runChatTurn, type ChatTurnEvent } from "./index" +import { MAX_STEP_ATTEMPTS } from "./budgets" +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 = { + orgId: "org_test" as TenantContext["orgId"], + userId: "user_test" as TenantContext["userId"], + roles: [], + authMode: "self_hosted", +} + +const MODEL: Model = CloudflareWorkersAI.configure({ + accountId: "test", + apiKey: "test", +}).model("@cf/test/model") + +/** A text delta, as the provider-neutral event the turn folds. */ +const textDelta = (text: string): LLMEvent => ({ type: "text-delta", id: "t1", text }) as LLMEvent + +const finish = (): LLMEvent => ({ type: "finish", reason: "stop" }) as LLMEvent + +const toolCall = (id: string, name: string): LLMEvent => + ({ type: "tool-call", id, name, input: {}, providerExecuted: false }) as LLMEvent + +/** + * Stub the model with a scripted event stream per step. + * + * `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 + * `./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. */ + | Failure + +/** Every request the turn issued, in order. Lets a test assert on tools, toolChoice and messages. */ +type RequestLog = Array + +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) => { + 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: partial.reason ?? "ProviderInternal" }, + message: "upstream exploded", + retryable: partial.retryable ?? partial.reason === undefined, + } as never + return Stream.concat(Stream.fromIterable(partial.events), Stream.fail(failure)) + }, + } + return { layer: Layer.succeed(LLMClient.Service, service as never), log, calls: () => step } +} + +const collect = ( + steps: ReadonlyArray, + 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(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) + +const terminal = (events: ReadonlyArray) => events.filter((event) => event.type === "turn-end") + +describe("runChatTurn", () => { + it("emits turn-start, the text, and exactly one turn-end", async () => { + const events = await Effect.runPromise(collectEvents([[textDelta("Hello"), finish()]])) + + assert.deepEqual(types(events), ["turn-start", "text-delta", "turn-end"]) + assert.lengthOf(terminal(events), 1) + }) + + it("coalesces adjacent text deltas into one event without losing any text", async () => { + const chunks = ["Check", "ing ", "the ", "traces", "."] + 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. + const deltas = events.filter((event) => event.type === "text-delta") + assert.lengthOf(deltas, 1) + assert.equal( + deltas.map((event) => (event.type === "text-delta" ? event.text : "")).join(""), + chunks.join(""), + ) + }) + + it("keeps text ahead of the tool calls it precedes", async () => { + const events = await Effect.runPromise( + collectEvents([ + [textDelta("Looking"), textDelta(" it up"), toolCall("c1", "create_alert_rule"), finish()], + ]), + ) + + // Batching must never reorder: the deltas live in one stream segment and the tool events in + // the concatenated one after it, so a slow batch cannot overtake the call it introduced. + assert.deepEqual(types(events), ["turn-start", "text-delta", "tool-call", "turn-end"]) + const delta = events.find((event) => event.type === "text-delta") + assert.equal(delta?.type === "text-delta" ? delta.text : undefined, "Looking it up") + }) + + 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. + // Both landed in the durable log; the SSE route stops at the first, so it only surfaced on + // the next reload. + assert.lengthOf(terminal(events), 1) + const end = terminal(events)[0] + assert.equal(end?.type === "turn-end" ? end.reason : undefined, "error") + }) + + it("dispatches NO tools when the stream fails after announcing one", async () => { + // 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( + collectEvents([ + { + events: [toolCall("c1", "find_errors"), textDelta("partial")], + fail: true, + reason: "Authentication", + }, + ]), + ) + + assert.isEmpty( + events.filter((event) => event.type === "tool-result"), + "a failed turn must not run tools past its terminal event", + ) + }) + + it("stops on an approval-gated tool with a proposal and no result", async () => { + 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) + assert.equal( + proposals[0]?.type === "tool-call" ? proposals[0].proposed : undefined, + true, + "a gated call is a proposal, not an execution", + ) + // Nothing fabricates an outcome: the model is never told the mutation happened. + assert.isEmpty(events.filter((event) => event.type === "tool-result")) + assert.lengthOf(terminal(events), 1) + }) + + it("stops without a second terminal event when the turn is superseded", async () => { + const events = await Effect.runPromise( + collectEvents([[textDelta("hi"), finish()]], { isCurrent: () => false }), + ) + + // An abort already recorded the terminal event on the session; emitting another here would + // double-close the turn in the durable log. + 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 `./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/loop/turn.ts b/apps/api/src/chat/loop/turn.ts new file mode 100644 index 000000000..cc765b377 --- /dev/null +++ b/apps/api/src/chat/loop/turn.ts @@ -0,0 +1,425 @@ +/** + * The chat turn loop: one submission, driven to completion. + * + * This module is the control flow and nothing else. Every policy it applies lives next door and is + * read here as a decision, not re-derived: + * + * - `./budgets.ts` — every ceiling: steps, attempts, backoff time, delegation fan-out, depth. + * - `./retry.ts` — which failures are worth another attempt, and how long to wait. + * - `./context.ts` — keeping the request inside the model's window. + * - `./delegate.ts` — handing a sub-question to a sub-agent, which re-enters this loop. + * - `./types.ts` — the input, the events, and one step's state. + * - `../tools.ts` — what the model may call. + * - `../agents.ts` — which agent this turn runs as, and what it is allowed to do. + * + * ## The shape of a step + * + * Stream the model → fold the events into an `LLMResponse` → settle the tool calls → build the next + * request → recurse. Recursion rather than a `while` because each step's output is a `Stream` that + * must be concatenated lazily: the next request cannot be built until the current step's tool + * results exist. + * + * ## Approvals are a real interrupt + * + * Flue's event stream had no human-in-the-loop primitive, so `apps/chat-flue/src/lib/approval.ts` + * swapped every mutating tool for one whose `execute` returned a `{ status: "proposed" }` marker + * without mutating — a propose-then-apply stub the web client re-ran through `POST /api/chat/apply`. + * Here the loop simply *stops* on a gated call: it emits a `tool-call` with `proposed: true` and + * ends the turn. Nothing fabricates a tool result, so the model is never told a mutation happened + * when it did not. `POST /api/chat/apply` is still how the approved mutation runs — the user's own + * action, authenticated as the user, which is where it belongs. + */ +import { evaluatePermission } from "@maple/domain/permission" +import { + LLM, + LLMEvent, + LLMResponse, + Message, + ToolResultPart, + toDefinitions, + ToolRuntime, + type LLMRequest, + type Tools, +} from "@maple/llm" +import { Duration, Effect, Stream } from "effect" +import { contextLimitOf, outputLimitOf, toLlmCallError } from "@/platform/Llm" +import { agentForSession, buildSystemPrompt, spawnableFor } from "../agents" +import { buildChatTools } from "../tools" +import { + hasStepBudget, + makeStepRetryBudget, + makeTaskBudget, + MAX_STEP_ATTEMPTS, + MAX_STEPS, + STEP_RETRY_BUDGET_MS, + TOOL_CONCURRENCY, +} from "./budgets" +import { isNearContextLimit, pruneToolResults } from "./context" +import { buildTaskTool } from "./delegate" +import { isRetryableStepFailure, stepRetryDelayMs } from "./retry" +import { + addUsage, + isCurrent, + tagged, + turnEnd, + type ChatTurnEvent, + type ChatTurnInput, + type StepState, +} from "./types" + +/** + * 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." + +/** + * How many text deltas are folded into one emitted event, and how long a partial batch waits. + * + * Roughly one animation frame. Every delta that leaves this stream becomes a durable SQLite row, an + * SSE frame and a React state commit, and the browser cannot show more than one update per frame + * anyway — so batching to that granularity costs no perceptible smoothness and removes most of the + * per-token work at all three layers. The size cap keeps a fast provider from letting a batch grow + * unboundedly within the window. + */ +const DELTA_BATCH_SIZE = 24 +const DELTA_BATCH_WINDOW = "16 millis" + +/** + * Run one submission to completion, streaming `ChatTurnEvent`s. + * + * The model is *streamed*, not `generate`d, so text deltas reach the session log — and through it + * the client — while the turn is still running. The raw `LLMEvent`s are folded into an + * `LLMResponse` on the way past so the assistant turn can be appended to the transcript verbatim + * for the next step. + */ +export const runChatTurn = (input: ChatTurnInput): Stream.Stream => + Stream.unwrap( + Effect.sync(() => { + 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: buildSystemPrompt(agent), + messages: [...input.messages], + tools: toDefinitions(tools), + }) + 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)) + }), + ) + +/** + * One assistant turn, then either settle its tool calls and recurse, or stop. + * + * Recursion rather than a loop because each step's output is a `Stream` that must be concatenated + * lazily: the next request cannot be built until the current step's tool results exist. + */ +const runStep = ( + input: ChatTurnInput, + tools: Tools, + request: LLMRequest, + 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 + // *second* terminal event after the error one — and, if that partial response carried tool + // calls, would dispatch them and recurse after the turn had already been declared over. + // 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))), + Stream.filter((event) => event.type === "text-delta" && event.text !== ""), + // One durable row, one SSE frame and one React commit per *token* is more fidelity than + // a screen can show. Coalescing into roughly one frame's worth of deltas is invisible + // to a reader and cuts all three by about an order of magnitude. Only text deltas are + // batched, and only against each other — `collected` still holds every raw event, and + // 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 => { + 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 + 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, + ), + ), + ), + ) + }), + ) + + const settleAndRecurse = Stream.unwrap( + Effect.gen(function* () { + if (failed) return Stream.empty + // Aborted between steps: the session already recorded the terminal event, so stop + // without emitting a second one. + if (!isCurrent(input)) return Stream.empty + + const response = LLMResponse.fromEvents(collected) + // A stream that neither failed nor assembled still ended the turn; say so, rather + // than leaving the log with no terminal event at all. + if (!response) return Stream.fromIterable([turnEnd(input, "stop")]) + + if (input.usage) addUsage(input.usage, response.usage) + + const calls = response.events + .filter(LLMEvent.is.toolCall) + .filter((call) => !call.providerExecuted) + + 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) => evaluatePermission(state.agent.permission, call.name) === "ask", + ) + if (gated) { + 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) => + tagged(input, { + type: "tool-call" as const, + messageId: input.messageId, + callId: call.id, + name: call.name, + input: call.input, + }), + ) + + // Announce first, settle second, as two stream segments. Emitting both together + // after `Effect.forEach` resolved meant a tool call only ever reached the log + // *already finished*, so the UI could never render one running — most of the point + // of streaming a turn that spends its time in tools. + const settled = Stream.unwrap( + Effect.gen(function* () { + const dispatched = yield* Effect.forEach( + calls, + (call) => + ToolRuntime.dispatch(tools, call).pipe( + Effect.map((result) => [call, result] as const), + ), + { concurrency: TOOL_CONCURRENCY }, + ) + + const results = dispatched.map(([call, outcome]) => + tagged(input, { + type: "tool-result" as const, + messageId: input.messageId, + callId: call.id, + output: outcome.result.value, + ...(outcome.result.type === "error" ? { isError: true } : {}), + }), + ) + + // Aborted while the tools were in flight: record what they returned so the + // transcript is not left with dangling calls, then stop. + if (!isCurrent(input)) return Stream.fromIterable(results) + + 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 = 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, { + ...state, + step: state.step + 1, + attempt: 0, + }), + ) + }), + ) + + return Stream.concat(Stream.fromIterable(announced), settled) + }), + ) + + return Stream.concat(live, settleAndRecurse) + }) diff --git a/apps/api/src/chat/loop/types.ts b/apps/api/src/chat/loop/types.ts new file mode 100644 index 000000000..fc5fbd1bc --- /dev/null +++ b/apps/api/src/chat/loop/types.ts @@ -0,0 +1,140 @@ +/** + * The vocabulary a turn speaks: what goes in, what comes out, and what one step carries. + * + * Separate from `turn.ts` so the loop's control flow reads as control flow. Everything here is data + * plus two small stamping helpers; nothing here decides anything. + */ +import type { ChatEvent, ChatTaskRef } from "@maple/domain/chat-session" +import type { Message, Model, Tools, Usage } from "@maple/llm" +import type { TenantContext } from "@/services/auth/tenant-context" +import type { AgentDefinition } from "../agents" +import type { StepRetryBudget, TaskBudget } from "./budgets" + +/** Distributive `Omit`, so each union member keeps its own shape. */ +type WithoutSeq = T extends unknown ? Omit : never + +/** + * Events this turn wants appended to the session log. `seq` is assigned by the Durable Object, + * which owns the ordering, so the loop emits everything without one. `user-message` is excluded: + * the session records the user turn at submission time, before the loop runs. + */ +export type ChatTurnEvent = WithoutSeq> + +export interface ChatTurnInput { + readonly sessionId: string + readonly tenant: TenantContext + readonly model: Model + /** The full transcript so far, oldest first, already including the new user message. */ + readonly messages: ReadonlyArray + readonly messageId: string + /** + * Investigate-mode sessions get a `submit_diagnosis` tool. It is supplied rather than built + * inside the loop because it needs `InvestigationService`, which would otherwise drag the + * service graph into the loop's imports. + */ + readonly extraTools?: Tools + /** + * Whether this turn still holds the session's turn slot. + * + * Checked between steps so an abort takes effect at the next boundary instead of only after the + * in-flight model call drains, and so a turn that has been superseded stops writing into a + * conversation that has moved on. Defaults to "always current" for callers with no session. + */ + 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 +} + +/** Where one step sits in the turn. `attempt` is 0-based and resets per step. */ +export 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 +} + +/** + * Running token total for one turn, accumulated across its steps. + * + * Mutable and shared rather than returned, because the one consumer — `submit_diagnosis` — is a + * *tool* invoked mid-turn, so there is no "after the turn" moment at which to hand it a total. + * In practice the diagnosis call is the last thing an investigation does, so this is the whole turn + * bar the final assistant message. Before this, `SubmitDiagnosisRequest` was built with no usage at + * all, so `InvestigationService`'s `if (env && (inputTokens || outputTokens))` was always false: + * `investigations.model` stayed null and Autumn was never metered for autonomous investigations, + * which the pre-`@maple/llm` workflow path did meter. + */ +export interface TurnUsage { + input: number + output: number + cacheRead: number +} + +export const makeTurnUsage = (): TurnUsage => ({ input: 0, output: 0, cacheRead: 0 }) + +export const addUsage = (total: TurnUsage, usage: Usage | undefined): void => { + total.input += usage?.inputTokens ?? 0 + total.output += usage?.outputTokens ?? 0 + total.cacheRead += usage?.cacheReadInputTokens ?? 0 +} + +/** + * 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. Typed structurally so this module does not depend on the + * whole turn input. + */ +export const tagged = (source: { readonly task?: ChatTaskRef }, event: E): E => + source.task === undefined ? event : { ...event, task: source.task } + +export const turnEnd = ( + input: { readonly messageId: string; readonly task?: ChatTaskRef }, + reason: Extract["reason"], + error?: string, +): 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. */ +export const isCurrent = (input: ChatTurnInput): boolean => input.isCurrent === undefined || input.isCurrent() 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/session.ts b/apps/api/src/chat/session.ts index a0c20997e..f003ba1af 100644 --- a/apps/api/src/chat/session.ts +++ b/apps/api/src/chat/session.ts @@ -3,7 +3,7 @@ * * This module is deliberately tiny and dependency-free: it is imported by `ai-triage-enqueue` and * `InvestigationService`, which are themselves reachable from the MCP tool registry, so anything - * heavy here would close an import cycle back through `chat/agent.ts`. + * heavy here would close an import cycle back through `chat/loop`. * * Starting a turn is now a single `beginTurn` call. Under Flue there were two very different paths * into the same conversation — the browser POSTed to `/agents/maple-chat/:id` on the chat-flue diff --git a/apps/api/src/chat/tools.ts b/apps/api/src/chat/tools.ts new file mode 100644 index 000000000..853e5145d --- /dev/null +++ b/apps/api/src/chat/tools.ts @@ -0,0 +1,106 @@ +/** + * What a chat turn is allowed to call. + * + * Kept out of `loop/` on purpose: the loop's job is to decide *when* to call a tool and what to do + * with the result, not to know which tools exist. Swapping the tool set — a read-only sub-agent, a + * mode with narrower reach — should not touch the control flow at all. + */ +import { investigationIdFromChatSessionId } 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 { Tool, ToolFailure, type Model, type Tools } from "@maple/llm" +import { Effect, Option, Schema } from "effect" +import { buildMapleTools, summarizeToolFailure, withRuntimeServices } from "@/mcp/tools/llm-tools" +import type { TenantContext } from "@/services/auth/tenant-context" +import type { TurnUsage } from "./loop/types" + +const decodeInvestigationIdOption = Schema.decodeUnknownOption(InvestigationId) + +/** + * The `submit_diagnosis` tool for an investigate-mode session (`":inv-"`). + * + * The agent calls it exactly once at the end of its autonomous pass and its arguments ARE the + * structured report — `AiTriageResult` directly, not the Valibot mirror `apps/chat-flue` had to + * keep in sync by hand. + * + * Deliberately not approval-gated: it is the structured-output channel, not a user-facing + * mutation. The investigation id and org ride from the session id, so the agent never chooses + * which investigation it writes. + * + * `submitDiagnosis` arrives as a callback rather than being resolved from `InvestigationService` + * here: that service is itself what starts an investigation's autonomous turn, so resolving it + * through the Effect requirements channel would make `InvestigationService` require itself. + */ +export type SubmitDiagnosis = ( + orgId: TenantContext["orgId"], + investigationId: InvestigationId, + request: SubmitDiagnosisRequest, + // eslint-disable-next-line @typescript-eslint/no-explicit-any +) => Effect.Effect + +export const buildSubmitDiagnosisTool = ( + sessionId: string, + tenant: TenantContext, + submitDiagnosis: SubmitDiagnosis, + usage: TurnUsage, + model: Model, +): Tools => { + const tools: Tools = {} + const rawId = investigationIdFromChatSessionId(sessionId) + if (!rawId) return tools + // `decodeUnknownSync` here turned a session id whose `inv-` suffix was not a UUID into a thrown + // defect on a user-supplied string. An unparseable id simply means this conversation is not an + // investigation, so it gets no `submit_diagnosis` tool. + const decoded = decodeInvestigationIdOption(rawId) + if (Option.isNone(decoded)) return tools + const investigationId = decoded.value + tools.submit_diagnosis = Tool.make({ + description: + "Record your structured diagnosis for THIS investigation. Call it exactly once, " + + "after you have gathered evidence, with your final assessment. It persists the report " + + "and renders it for the user. After calling it, stop unless the user asks a follow-up.", + parameters: AiTriageResult, + success: Schema.String, + execute: (report) => + withRuntimeServices( + submitDiagnosis( + tenant.orgId, + investigationId, + new SubmitDiagnosisRequest({ + report, + model: String(model.id), + inputTokens: usage.input, + outputTokens: usage.output, + }), + ).pipe( + Effect.as("Diagnosis recorded."), + // Named failures only. `catchCause` + `String(cause)` fed the model a rendered + // Effect cause — stack frames, and connection details out of a DatabaseError. + Effect.catchCause((cause) => + Effect.fail( + new ToolFailure({ + message: `submit_diagnosis failed: ${summarizeToolFailure(cause)}`, + }), + ), + ), + ), + ), + }) + return tools +} + +/** + * 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. + */ +export 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", + }) 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..48255bd8f 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("./loop")) + 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. * @@ -113,12 +194,14 @@ const toLlmMessages = (history: ReadonlyArray): ReadonlyArray => { - const [{ MainLive }, { layerPg }, { layerLlm, resolveTriageModel }, agent] = await Promise.all([ - import("../app"), - import("../platform/DatabasePgLive"), - import("../platform/Llm"), - import("./agent"), - ]) + const [{ MainLive }, { layerPg }, { layerLlm, resolveTriageModel }, loop, { buildSubmitDiagnosisTool }] = + await Promise.all([ + import("../app"), + import("../platform/DatabasePgLive"), + import("../platform/Llm"), + import("./loop"), + import("./tools"), + ]) const { InvestigationService } = await import("@/services/errors/InvestigationService") const runtime = ManagedRuntime.make( @@ -143,8 +226,8 @@ export const runChatSessionTurn = async (input: RunChatSessionTurnInput): Promis }) // Shared with the turn so `submit_diagnosis` can report what the investigation cost. See // `TurnUsage` — the tool is invoked mid-turn, so there is no later moment to hand it a total. - const usage = agent.makeTurnUsage() - const extraTools = agent.buildSubmitDiagnosisTool( + const usage = loop.makeTurnUsage() + const extraTools = buildSubmitDiagnosisTool( input.sessionId, tenant, investigations.submitDiagnosis, @@ -152,23 +235,32 @@ export const runChatSessionTurn = async (input: RunChatSessionTurnInput): Promis model, ) - yield* agent + yield* loop .runChatTurn({ 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/dispatcher.ts b/apps/api/src/mcp/dispatcher.ts index 74308fb0e..795fa6324 100644 --- a/apps/api/src/mcp/dispatcher.ts +++ b/apps/api/src/mcp/dispatcher.ts @@ -26,8 +26,8 @@ const toDecodeErrorMessage = (definition: MapleToolDefinition, error: unknown): /** * Built on first use, not at module scope. * - * `apps/api/src/chat/agent.ts` imports this module and is itself reachable from the tool registry's - * own import graph (registry -> a tool -> issue-hub/ai-triage-enqueue -> chat/session -> chat/agent + * `apps/api/src/chat/tools.ts` imports this module and is itself reachable from the tool registry's + * own import graph (registry -> a tool -> issue-hub/ai-triage-enqueue -> chat/session -> chat/tools * -> here). Computing the descriptors eagerly meant that whichever module the bundler happened to * evaluate first could observe `mapleToolDefinitions` as `undefined`. Deferring removes the * ordering dependency entirely rather than papering over one edge of the cycle. 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..b1b63aaf8 --- /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/tools.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..c291bf65c 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/loop" 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("|")}