diff options
| author | Adam Malczewski <[email protected]> | 2026-06-11 12:23:06 +0900 |
|---|---|---|
| committer | Adam Malczewski <[email protected]> | 2026-06-11 12:23:06 +0900 |
| commit | c2b4c05d91fa88b8d02c055a0e15c22abd8e21f3 (patch) | |
| tree | 3f7c2feddbe697a79abd952bb80ed0e01dac0a7a /packages/session-orchestrator/src | |
| parent | f6b45507210e04e9884256b0132900640de4334b (diff) | |
| download | dispatch-c2b4c05d91fa88b8d02c055a0e15c22abd8e21f3.tar.gz dispatch-c2b4c05d91fa88b8d02c055a0e15c22abd8e21f3.zip | |
feat(cache-warming): per-conversation prompt-cache warming + warm() service
Backend-driven warming targeting whatever provider a conversation uses (incl. the
external Claude provider-anthropic). Core engine + on/off + last-cache-% done;
interval-as-view-control pending a ui-contract NumberField (surface-system gap).
Mechanism:
- kernel: expose HostAPI.emit (typed bus event emit; counterpart of on)
- session-orchestrator: turnStarted/turnSettled event hooks (conversationId/cwd/model);
warm() service (cacheWarmHandle) reusing the real-turn assembly (byte-identical prefix,
provider-agnostic), refuses mid-turn, never persists/emits, returns Usage
- cache-warming (new ext): per-conversation timers (arm on settle, cancel on start,
in-flight invalidation), calls warm(), pct=round(clamp(cacheRead/input,0,1)*100),
persists {enabled,intervalMs} (default on/240s), registers a controls surface
- host-bin: register cache-warming; transport-http: HostAPI stub +emit (fan-out)
Honors old-code invariants. 760 vitest + 109 bun = 869 tests; tsc -b EXIT 0; biome clean.
Diffstat (limited to 'packages/session-orchestrator/src')
| -rw-r--r-- | packages/session-orchestrator/src/extension.ts | 31 | ||||
| -rw-r--r-- | packages/session-orchestrator/src/index.ts | 8 | ||||
| -rw-r--r-- | packages/session-orchestrator/src/orchestrator.test.ts | 296 | ||||
| -rw-r--r-- | packages/session-orchestrator/src/orchestrator.ts | 219 |
4 files changed, 486 insertions, 68 deletions
diff --git a/packages/session-orchestrator/src/extension.ts b/packages/session-orchestrator/src/extension.ts index 697eb4a..12d387c 100644 --- a/packages/session-orchestrator/src/extension.ts +++ b/packages/session-orchestrator/src/extension.ts @@ -3,8 +3,9 @@ import { credentialStoreHandle } from "@dispatch/credential-store"; import type { Extension, HostAPI, Manifest } from "@dispatch/kernel"; import { runTurn } from "@dispatch/kernel"; import { + cacheWarmHandle, createSessionOrchestrator, - type SessionOrchestrator, + createWarmService, sessionOrchestratorHandle, } from "./orchestrator.js"; import { selectFirstProvider } from "./pure.js"; @@ -19,14 +20,15 @@ export const manifest: Manifest = { dependsOn: ["conversation-store", "credential-store"], activation: "eager", contributes: { - services: ["session-orchestrator/orchestrator"], + services: ["session-orchestrator/orchestrator", "session-orchestrator/warm"], + hooks: ["session-orchestrator/turn-started", "session-orchestrator/turn-settled"], }, }; export function activate(host: HostAPI): void { const conversationStore = host.getService(conversationStoreHandle); - const orchestrator: SessionOrchestrator = createSessionOrchestrator({ + const { orchestrator, activeConversations } = createSessionOrchestrator({ conversationStore, resolveProvider: () => selectFirstProvider(host.getProviders()), resolveTools: () => [...host.getTools().values()], @@ -41,9 +43,32 @@ export function activate(host: HostAPI): void { runTurn, logger: host.logger, now: () => Date.now(), + emit: (hook, payload) => host.emit(hook, payload), }); host.provideService(sessionOrchestratorHandle, orchestrator); + + const warmService = createWarmService( + { + conversationStore, + resolveProvider: () => selectFirstProvider(host.getProviders()), + resolveTools: () => [...host.getTools().values()], + resolveModel: (modelName: string) => { + const store = host.getService(credentialStoreHandle); + const r = store.resolve(modelName); + if (r === undefined) return undefined; + const provider = host.getProviders().get(r.providerId); + return provider ? { provider, model: r.model } : undefined; + }, + applyToolsFilter: (assembly) => host.applyFilters(toolsFilter, assembly), + runTurn, + logger: host.logger, + now: () => Date.now(), + }, + activeConversations, + ); + + host.provideService(cacheWarmHandle, warmService); } export const extension: Extension = { diff --git a/packages/session-orchestrator/src/index.ts b/packages/session-orchestrator/src/index.ts index 071b616..37ae5ce 100644 --- a/packages/session-orchestrator/src/index.ts +++ b/packages/session-orchestrator/src/index.ts @@ -1,9 +1,17 @@ export { extension, manifest } from "./extension.js"; export { + cacheWarmHandle, createSessionOrchestrator, + createWarmService, type SessionOrchestrator, + type SessionOrchestratorBundle, type SessionOrchestratorDeps, sessionOrchestratorHandle, + type TurnLifecyclePayload, + turnSettled, + turnStarted, + type WarmResult, + type WarmService, } from "./orchestrator.js"; export { buildUserMessage, diff --git a/packages/session-orchestrator/src/orchestrator.test.ts b/packages/session-orchestrator/src/orchestrator.test.ts index dcaad7d..5d512ea 100644 --- a/packages/session-orchestrator/src/orchestrator.test.ts +++ b/packages/session-orchestrator/src/orchestrator.test.ts @@ -2,6 +2,7 @@ import type { ConversationStore } from "@dispatch/conversation-store"; import type { AgentEvent, ChatMessage, + EventHookDescriptor, ProviderContract, ProviderEvent, RunTurnInput, @@ -12,7 +13,11 @@ import type { } from "@dispatch/kernel"; import { runTurn } from "@dispatch/kernel"; import { describe, expect, it } from "vitest"; -import { createSessionOrchestrator } from "./orchestrator.js"; +import { + createSessionOrchestrator, + createWarmService, + type TurnLifecyclePayload, +} from "./orchestrator.js"; import type { ToolAssembly } from "./tools-filter.js"; function createInMemoryStore(): ConversationStore & { @@ -104,7 +109,7 @@ describe("handleMessage integration", () => { ], ]); - const orchestrator = createSessionOrchestrator({ + const { orchestrator } = createSessionOrchestrator({ conversationStore: store, resolveProvider: () => provider, resolveTools: () => [], @@ -156,7 +161,7 @@ describe("handleMessage integration", () => { }, }; - const orchestrator = createSessionOrchestrator({ + const { orchestrator } = createSessionOrchestrator({ conversationStore: store, resolveProvider: () => provider, resolveTools: () => [], @@ -203,7 +208,7 @@ describe("handleMessage integration", () => { ], ]); - const orchestrator = createSessionOrchestrator({ + const { orchestrator } = createSessionOrchestrator({ conversationStore: store, resolveProvider: () => provider, resolveTools: () => [], @@ -233,7 +238,7 @@ describe("handleMessage integration", () => { ], ]); - const orchestrator = createSessionOrchestrator({ + const { orchestrator } = createSessionOrchestrator({ conversationStore: store, resolveProvider: () => provider, resolveTools: () => [], @@ -282,7 +287,7 @@ describe("handleMessage model resolution", () => { const fallbackProvider: ProviderContract = { id: "fallback", stream: async function* () {} }; const { captured, captureRunTurn } = createCapturingRunTurn(); - const orchestrator = createSessionOrchestrator({ + const { orchestrator } = createSessionOrchestrator({ conversationStore: store, resolveProvider: () => fallbackProvider, resolveTools: () => [], @@ -314,7 +319,7 @@ describe("handleMessage model resolution", () => { const { captured, captureRunTurn } = createCapturingRunTurn(); const events: AgentEvent[] = []; - const orchestrator = createSessionOrchestrator({ + const { orchestrator } = createSessionOrchestrator({ conversationStore: store, resolveProvider: () => fallbackProvider, resolveTools: () => [], @@ -345,7 +350,7 @@ describe("handleMessage model resolution", () => { const fallbackProvider: ProviderContract = { id: "fallback", stream: async function* () {} }; const { captured, captureRunTurn } = createCapturingRunTurn(); - const orchestrator = createSessionOrchestrator({ + const { orchestrator } = createSessionOrchestrator({ conversationStore: store, resolveProvider: () => fallbackProvider, resolveTools: () => [], @@ -373,7 +378,7 @@ describe("handleMessage model resolution", () => { const provider: ProviderContract = { id: "p", stream: async function* () {} }; const { captured, captureRunTurn } = createCapturingRunTurn(); - const orchestrator = createSessionOrchestrator({ + const { orchestrator } = createSessionOrchestrator({ conversationStore: store, resolveProvider: () => provider, resolveTools: () => [], @@ -407,7 +412,7 @@ describe("handleMessage model resolution", () => { const { captured, captureRunTurn } = createCapturingRunTurn(); const fakeNow = () => 42; - const orchestrator = createSessionOrchestrator({ + const { orchestrator } = createSessionOrchestrator({ conversationStore: store, resolveProvider: () => provider, resolveTools: () => [], @@ -432,7 +437,7 @@ describe("handleMessage model resolution", () => { const provider: ProviderContract = { id: "p", stream: async function* () {} }; const { captured, captureRunTurn } = createCapturingRunTurn(); - const orchestrator = createSessionOrchestrator({ + const { orchestrator } = createSessionOrchestrator({ conversationStore: store, resolveProvider: () => provider, resolveTools: () => [], @@ -461,7 +466,7 @@ describe("turn-sealed event", () => { ], ]); - const orchestrator = createSessionOrchestrator({ + const { orchestrator } = createSessionOrchestrator({ conversationStore: store, resolveProvider: () => provider, resolveTools: () => [], @@ -514,7 +519,7 @@ describe("turn-sealed event", () => { }, }; - const orchestrator = createSessionOrchestrator({ + const { orchestrator } = createSessionOrchestrator({ conversationStore: wrappedStore, resolveProvider: () => provider, resolveTools: () => [], @@ -561,7 +566,7 @@ describe("turn-sealed event", () => { }, }; - const orchestrator = createSessionOrchestrator({ + const { orchestrator } = createSessionOrchestrator({ conversationStore: failingStore, resolveProvider: () => provider, resolveTools: () => [], @@ -595,7 +600,7 @@ describe("turn metrics persistence", () => { ], ]); - const orchestrator = createSessionOrchestrator({ + const { orchestrator } = createSessionOrchestrator({ conversationStore: store, resolveProvider: () => provider, resolveTools: () => [], @@ -655,7 +660,7 @@ describe("turn metrics persistence", () => { }, }; - const orchestrator = createSessionOrchestrator({ + const { orchestrator } = createSessionOrchestrator({ conversationStore: store, resolveProvider: () => provider, resolveTools: () => [tool], @@ -714,7 +719,7 @@ describe("turn metrics persistence", () => { }, }; - const orchestrator = createSessionOrchestrator({ + const { orchestrator } = createSessionOrchestrator({ conversationStore: store, resolveProvider: () => provider, resolveTools: () => [], @@ -781,7 +786,7 @@ describe("turn metrics persistence", () => { }, }; - const orchestrator = createSessionOrchestrator({ + const { orchestrator } = createSessionOrchestrator({ conversationStore: store, resolveProvider: () => provider, resolveTools: () => [tool], @@ -835,7 +840,7 @@ describe("turn metrics persistence", () => { }, }; - const orchestrator = createSessionOrchestrator({ + const { orchestrator } = createSessionOrchestrator({ conversationStore: failingMetricsStore, resolveProvider: () => provider, resolveTools: () => [], @@ -874,7 +879,7 @@ describe("tools filter", () => { return Promise.resolve({ ...assembly, tools: [toolB] }); }; - const orchestrator = createSessionOrchestrator({ + const { orchestrator } = createSessionOrchestrator({ conversationStore: store, resolveProvider: () => provider, resolveTools: () => [toolA], @@ -902,7 +907,7 @@ describe("tools filter", () => { const toolA = createFakeTool("tool-a", async () => ({ content: "a" })); const toolB = createFakeTool("tool-b", async () => ({ content: "b" })); - const orchestrator = createSessionOrchestrator({ + const { orchestrator } = createSessionOrchestrator({ conversationStore: store, resolveProvider: () => provider, resolveTools: () => [toolA, toolB], @@ -933,7 +938,7 @@ describe("tools filter", () => { return Promise.resolve(assembly); }; - const orchestrator = createSessionOrchestrator({ + const { orchestrator } = createSessionOrchestrator({ conversationStore: store, resolveProvider: () => provider, resolveTools: () => [], @@ -964,3 +969,250 @@ function createCounterNow(): { now: () => number; tick: (ms: number) => void } { }, }; } + +describe("lifecycle event hooks", () => { + it("emits turnStarted before and turnSettled after a turn", async () => { + const store = createInMemoryStore(); + const provider = createFakeProvider([ + [ + { type: "text-delta", delta: "ok" }, + { type: "finish", reason: "stop" }, + ], + ]); + + const emitted: Array<{ hook: string; payload: TurnLifecyclePayload; order: number }> = []; + let order = 0; + + const fakeEmit = <TPayload>(hook: EventHookDescriptor<TPayload>, payload: TPayload): void => { + emitted.push({ hook: hook.id, payload: payload as TurnLifecyclePayload, order: order++ }); + }; + + const { orchestrator } = createSessionOrchestrator({ + conversationStore: store, + resolveProvider: () => provider, + resolveTools: () => [], + applyToolsFilter: identityApplyToolsFilter, + runTurn, + emit: fakeEmit, + }); + + await orchestrator.handleMessage({ + conversationId: "conv-lifecycle", + text: "test", + onEvent: () => {}, + cwd: "/work", + modelName: "mymodel", + }); + + expect(emitted).toHaveLength(2); + expect(emitted[0]?.hook).toBe("session-orchestrator/turn-started"); + expect(emitted[0]?.payload.conversationId).toBe("conv-lifecycle"); + expect(emitted[0]?.payload.cwd).toBe("/work"); + expect(emitted[0]?.payload.modelName).toBe("mymodel"); + expect(emitted[0]?.order).toBe(0); + + expect(emitted[1]?.hook).toBe("session-orchestrator/turn-settled"); + expect(emitted[1]?.payload.conversationId).toBe("conv-lifecycle"); + expect(emitted[1]?.payload.cwd).toBe("/work"); + expect(emitted[1]?.payload.modelName).toBe("mymodel"); + expect(emitted[1]?.order).toBe(1); + }); +}); + +describe("warm service", () => { + it("warm reuses the assembled tools + full history and appends the probe turn", async () => { + const store = createInMemoryStore(); + const existingMsg: ChatMessage = { + role: "user", + chunks: [{ type: "text", text: "existing" }], + }; + const assistantMsg: ChatMessage = { + role: "assistant", + chunks: [{ type: "text", text: "reply" }], + }; + await store.append("conv-warm-reuse", [existingMsg, assistantMsg]); + + let capturedMessages: readonly ChatMessage[] | undefined; + let capturedTools: readonly ToolContract[] | undefined; + let _capturedOpts: unknown; + + const toolA = createFakeTool("tool-a", async () => ({ content: "a" })); + + const provider: ProviderContract = { + id: "warm-provider", + stream(messages, tools, opts) { + capturedMessages = messages; + capturedTools = tools; + _capturedOpts = opts; + return (async function* () { + yield { + type: "usage", + usage: { inputTokens: 100, outputTokens: 5, cacheReadTokens: 80, cacheWriteTokens: 20 }, + } as ProviderEvent; + yield { type: "finish", reason: "stop" } as ProviderEvent; + })(); + }, + }; + + const deps = { + conversationStore: store, + resolveProvider: () => provider, + resolveTools: () => [toolA], + applyToolsFilter: identityApplyToolsFilter, + runTurn, + }; + + const { activeConversations } = createSessionOrchestrator(deps); + const warmService = createWarmService(deps, activeConversations); + + const result = await warmService.warm("conv-warm-reuse", { cwd: "/test" }); + + expect(capturedMessages).toBeDefined(); + expect(capturedMessages).toHaveLength(3); + expect(capturedMessages?.[0]?.chunks[0]).toEqual({ type: "text", text: "existing" }); + expect(capturedMessages?.[1]?.chunks[0]).toEqual({ type: "text", text: "reply" }); + expect(capturedMessages?.[2]?.role).toBe("user"); + expect((capturedMessages?.[2]?.chunks[0] as { type: "text"; text: string }).text).toBe( + "reply with just a .", + ); + + expect(capturedTools).toHaveLength(1); + expect(capturedTools?.[0]?.name).toBe("tool-a"); + + if ("inputTokens" in result) { + expect(result.inputTokens).toBe(100); + expect(result.cacheReadTokens).toBe(80); + } + }); + + it("warm refuses while the conversation is generating", async () => { + const store = createInMemoryStore(); + let resolveRunTurn: (() => void) | undefined; + const runTurnBlocker = new Promise<void>((resolve) => { + resolveRunTurn = resolve; + }); + + const provider: ProviderContract = { + id: "p", + stream: async function* () { + yield { type: "text-delta", delta: "slow" } as ProviderEvent; + yield { type: "finish", reason: "stop" } as ProviderEvent; + }, + }; + + const blockingRunTurn = async (_input: RunTurnInput): Promise<RunTurnResult> => { + await runTurnBlocker; + return { + messages: [{ role: "assistant", chunks: [{ type: "text", text: "done" }] }], + usage: { inputTokens: 1, outputTokens: 1 }, + finishReason: "stop", + }; + }; + + const deps = { + conversationStore: store, + resolveProvider: () => provider, + resolveTools: () => [], + applyToolsFilter: identityApplyToolsFilter, + runTurn: blockingRunTurn, + }; + + const { orchestrator, activeConversations } = createSessionOrchestrator(deps); + const warmService = createWarmService(deps, activeConversations); + + const turnPromise = orchestrator.handleMessage({ + conversationId: "conv-blocking", + text: "test", + onEvent: () => {}, + }); + + const warmResult = await warmService.warm("conv-blocking"); + expect(warmResult).toEqual({ error: "conversation is generating" }); + + resolveRunTurn?.(); + await turnPromise; + }); + + it("warm never persists (no append) and emits no AgentEvents", async () => { + const store = createInMemoryStore(); + const existingMsg: ChatMessage = { + role: "user", + chunks: [{ type: "text", text: "existing" }], + }; + await store.append("conv-no-persist", [existingMsg]); + + const provider: ProviderContract = { + id: "p", + stream: async function* () { + yield { + type: "usage", + usage: { inputTokens: 10, outputTokens: 2 }, + } as ProviderEvent; + yield { type: "finish", reason: "stop" } as ProviderEvent; + }, + }; + + const deps = { + conversationStore: store, + resolveProvider: () => provider, + resolveTools: () => [], + applyToolsFilter: identityApplyToolsFilter, + runTurn, + }; + + const { activeConversations } = createSessionOrchestrator(deps); + const warmService = createWarmService(deps, activeConversations); + + const sizeBefore = store.data.get("conv-no-persist")?.length; + + await warmService.warm("conv-no-persist"); + + const sizeAfter = store.data.get("conv-no-persist")?.length; + expect(sizeAfter).toBe(sizeBefore); + }); + + it("warm returns provider usage (input + cacheReadTokens)", async () => { + const store = createInMemoryStore(); + const existingMsg: ChatMessage = { + role: "user", + chunks: [{ type: "text", text: "existing" }], + }; + await store.append("conv-usage", [existingMsg]); + + const provider: ProviderContract = { + id: "p", + stream: async function* () { + yield { + type: "usage", + usage: { + inputTokens: 500, + outputTokens: 3, + cacheReadTokens: 400, + cacheWriteTokens: 100, + }, + } as ProviderEvent; + yield { type: "finish", reason: "stop" } as ProviderEvent; + }, + }; + + const deps = { + conversationStore: store, + resolveProvider: () => provider, + resolveTools: () => [], + applyToolsFilter: identityApplyToolsFilter, + runTurn, + }; + + const { activeConversations } = createSessionOrchestrator(deps); + const warmService = createWarmService(deps, activeConversations); + + const result = await warmService.warm("conv-usage"); + + expect(result).toEqual({ + inputTokens: 500, + outputTokens: 3, + cacheReadTokens: 400, + cacheWriteTokens: 100, + }); + }); +}); diff --git a/packages/session-orchestrator/src/orchestrator.ts b/packages/session-orchestrator/src/orchestrator.ts index 2d1bbf5..c39bc06 100644 --- a/packages/session-orchestrator/src/orchestrator.ts +++ b/packages/session-orchestrator/src/orchestrator.ts @@ -2,19 +2,59 @@ import type { ConversationStore } from "@dispatch/conversation-store"; import type { AgentEvent, ChatMessage, + EventHookDescriptor, Logger, ProviderContract, + ProviderEvent, ProviderStreamOptions, RunTurnInput, RunTurnResult, ToolContract, ToolDispatchPolicy, + UsageEvent, } from "@dispatch/kernel"; -import { defineService } from "@dispatch/kernel"; +import { defineEventHook, defineService, type ServiceHandle } from "@dispatch/kernel"; import { createMetricsAccumulator } from "./metrics.js"; import { buildUserMessage, defaultDispatchPolicy, generateTurnId } from "./pure.js"; import type { ToolAssembly } from "./tools-filter.js"; +// --- Lifecycle event hooks --- + +/** Context carried on turn-lifecycle events, enough to replicate the turn's request prefix. */ +export interface TurnLifecyclePayload { + readonly conversationId: string; + readonly cwd?: string; + readonly modelName?: string; +} + +/** Fired when a turn STARTS driving a conversation (consumers cancel warming timers). */ +export const turnStarted: EventHookDescriptor<TurnLifecyclePayload> = + defineEventHook<TurnLifecyclePayload>("session-orchestrator/turn-started"); + +/** Fired when a turn SETTLES (sealed) for a conversation (consumers arm warming timers). */ +export const turnSettled: EventHookDescriptor<TurnLifecyclePayload> = + defineEventHook<TurnLifecyclePayload>("session-orchestrator/turn-settled"); + +// --- Warm service --- + +export interface WarmResult { + readonly inputTokens: number; + readonly outputTokens: number; + readonly cacheReadTokens: number; + readonly cacheWriteTokens: number; +} + +export interface WarmService { + readonly warm: ( + conversationId: string, + opts?: { readonly cwd?: string; readonly modelName?: string }, + ) => Promise<WarmResult | { readonly error: string }>; +} + +export const cacheWarmHandle: ServiceHandle<WarmService> = defineService<WarmService>( + "session-orchestrator/warm", +); + export interface SessionOrchestrator { handleMessage(input: { conversationId: string; @@ -45,28 +85,129 @@ export interface SessionOrchestratorDeps { readonly logger?: Logger; /** Injected monotonic-ish clock (ms) forwarded to RunTurnInput for timing events. */ readonly now?: () => number; + /** Emit a lifecycle event hook to subscribers. Injected from host. */ + readonly emit?: <TPayload>(hook: EventHookDescriptor<TPayload>, payload: TPayload) => void; } -export function createSessionOrchestrator(deps: SessionOrchestratorDeps): SessionOrchestrator { - return { +export interface SessionOrchestratorBundle { + readonly orchestrator: SessionOrchestrator; + /** The shared active-conversations set, for use by createWarmService. */ + readonly activeConversations: ReadonlySet<string>; +} + +export function createSessionOrchestrator( + deps: SessionOrchestratorDeps, +): SessionOrchestratorBundle { + const activeConversations = new Set<string>(); + + const orchestrator: SessionOrchestrator = { async handleMessage({ conversationId, text, onEvent, signal, modelName, cwd }) { + const payload: TurnLifecyclePayload = { + conversationId, + ...(cwd !== undefined ? { cwd } : {}), + ...(modelName !== undefined ? { modelName } : {}), + }; + deps.emit?.(turnStarted, payload); + activeConversations.add(conversationId); + + try { + const history = await deps.conversationStore.load(conversationId); + const userMsg = buildUserMessage(text); + const turnId = generateTurnId(); + + let provider: ProviderContract; + let modelOverride: string | undefined; + + if (modelName !== undefined && deps.resolveModel !== undefined) { + const resolved = deps.resolveModel(modelName); + if (resolved === undefined) { + onEvent({ + type: "error", + conversationId, + turnId, + message: `unknown model: ${modelName}`, + }); + return; + } + provider = resolved.provider; + modelOverride = resolved.model; + } else { + provider = deps.resolveProvider(); + } + + const baseTools = deps.resolveTools(); + const assembled = await deps.applyToolsFilter({ + tools: baseTools, + conversationId, + ...(cwd !== undefined ? { cwd } : {}), + }); + const dispatch = deps.resolveDispatch?.() ?? defaultDispatchPolicy(); + const turnLogger = deps.logger?.child({ conversationId, turnId }); + const metrics = createMetricsAccumulator(); + + const emitAndAccumulate = (event: AgentEvent): void => { + metrics.ingest(event); + onEvent(event); + }; + + const opts: RunTurnInput = { + provider, + messages: [...history, userMsg], + tools: assembled.tools, + dispatch, + emit: emitAndAccumulate, + conversationId, + turnId, + ...(modelOverride !== undefined + ? { providerOpts: { model: modelOverride } satisfies ProviderStreamOptions } + : {}), + ...(turnLogger !== undefined ? { logger: turnLogger } : {}), + ...(signal !== undefined ? { signal } : {}), + ...(cwd !== undefined ? { cwd } : {}), + ...(deps.now !== undefined ? { now: deps.now } : {}), + }; + + const result = await deps.runTurn(opts); + + const toPersist: ChatMessage[] = [userMsg, ...result.messages]; + await deps.conversationStore.append(conversationId, toPersist); + + const turnMetrics = metrics.build(turnId); + await deps.conversationStore.appendMetrics(conversationId, turnMetrics); + + onEvent({ type: "turn-sealed", conversationId, turnId }); + } finally { + activeConversations.delete(conversationId); + deps.emit?.(turnSettled, payload); + } + }, + }; + + return { orchestrator, activeConversations }; +} + +export function createWarmService( + deps: SessionOrchestratorDeps, + activeConversations: ReadonlySet<string>, +): WarmService { + return { + async warm(conversationId, opts) { + if (activeConversations.has(conversationId)) { + return { error: "conversation is generating" }; + } + const history = await deps.conversationStore.load(conversationId); - const userMsg = buildUserMessage(text); - const turnId = generateTurnId(); + if (history.length === 0) { + return { error: "no history" }; + } let provider: ProviderContract; let modelOverride: string | undefined; - if (modelName !== undefined && deps.resolveModel !== undefined) { - const resolved = deps.resolveModel(modelName); + if (opts?.modelName !== undefined && deps.resolveModel !== undefined) { + const resolved = deps.resolveModel(opts.modelName); if (resolved === undefined) { - onEvent({ - type: "error", - conversationId, - turnId, - message: `unknown model: ${modelName}`, - }); - return; + return { error: `unknown model: ${opts.modelName}` }; } provider = resolved.provider; modelOverride = resolved.model; @@ -75,46 +216,38 @@ export function createSessionOrchestrator(deps: SessionOrchestratorDeps): Sessio } const baseTools = deps.resolveTools(); + const cwd = opts?.cwd; const assembled = await deps.applyToolsFilter({ tools: baseTools, conversationId, ...(cwd !== undefined ? { cwd } : {}), }); - const dispatch = deps.resolveDispatch?.() ?? defaultDispatchPolicy(); - const turnLogger = deps.logger?.child({ conversationId, turnId }); - const metrics = createMetricsAccumulator(); - const emitAndAccumulate = (event: AgentEvent): void => { - metrics.ingest(event); - onEvent(event); + const probeMsg: ChatMessage = { + role: "user", + chunks: [{ type: "text", text: "reply with just a ." }], }; + const messages = [...history, probeMsg]; - const opts: RunTurnInput = { - provider, - messages: [...history, userMsg], - tools: assembled.tools, - dispatch, - emit: emitAndAccumulate, - conversationId, - turnId, - ...(modelOverride !== undefined - ? { providerOpts: { model: modelOverride } satisfies ProviderStreamOptions } - : {}), - ...(turnLogger !== undefined ? { logger: turnLogger } : {}), - ...(signal !== undefined ? { signal } : {}), - ...(cwd !== undefined ? { cwd } : {}), - ...(deps.now !== undefined ? { now: deps.now } : {}), - }; - - const result = await deps.runTurn(opts); + const providerOpts: ProviderStreamOptions | undefined = + modelOverride !== undefined ? { model: modelOverride, maxTokens: 1 } : { maxTokens: 1 }; - const toPersist: ChatMessage[] = [userMsg, ...result.messages]; - await deps.conversationStore.append(conversationId, toPersist); + let inputTokens = 0; + let outputTokens = 0; + let cacheReadTokens = 0; + let cacheWriteTokens = 0; - const turnMetrics = metrics.build(turnId); - await deps.conversationStore.appendMetrics(conversationId, turnMetrics); + for await (const event of provider.stream(messages, assembled.tools, providerOpts)) { + if ((event as ProviderEvent).type === "usage") { + const usageEvent = event as UsageEvent; + inputTokens = usageEvent.usage.inputTokens; + outputTokens = usageEvent.usage.outputTokens; + cacheReadTokens = usageEvent.usage.cacheReadTokens ?? 0; + cacheWriteTokens = usageEvent.usage.cacheWriteTokens ?? 0; + } + } - onEvent({ type: "turn-sealed", conversationId, turnId }); + return { inputTokens, outputTokens, cacheReadTokens, cacheWriteTokens }; }, }; } |
