diff --git a/src/local-agent-adapters.ts b/src/local-agent-adapters.ts index ce91c90f..ab5f5e5c 100644 --- a/src/local-agent-adapters.ts +++ b/src/local-agent-adapters.ts @@ -19,11 +19,13 @@ import { ProviderSchemaUnsupportedError, type LocalAgentRunInput, type LocalAgentRunResult, + type LocalAgentObserver, + type LocalAgentUsageSnapshot, } from "./local-agent-runtime.js"; export interface LocalAgentAdapter { readonly provider: LocalAgentProvider; - run(input: LocalAgentRunInput): Promise; + run(input: LocalAgentRunInput, observer?: LocalAgentObserver): Promise; } const ACP_COMMANDS: Record<"cursor" | "copilot", [string, ...string[]]> = { @@ -35,8 +37,9 @@ const PI_AGENT_TIMEOUT_MS = 120_000; export async function runLocalAgentProvider( provider: LocalAgentProvider, input: LocalAgentRunInput, + observer?: LocalAgentObserver, ): Promise { - const result = await runLocalAgentProviderResult(provider, input); + const result = await runLocalAgentProviderResult(provider, input, observer); if (result.isErr()) throw result.error; return result.value; } @@ -44,9 +47,10 @@ export async function runLocalAgentProvider( export async function runLocalAgentProviderResult( provider: LocalAgentProvider, input: LocalAgentRunInput, + observer?: LocalAgentObserver, ): Promise> { return Result.tryPromise({ - try: () => createLocalAgentAdapter(provider).run(input), + try: () => createLocalAgentAdapter(provider).run(input, observer), catch: (cause) => classifyAgentProviderError(provider, cause), }); } @@ -70,16 +74,16 @@ export function createLocalAgentAdapter(provider: LocalAgentProvider): LocalAgen class CodexLocalAgentAdapter implements LocalAgentAdapter { readonly provider = "codex" as const; - async run(input: LocalAgentRunInput): Promise { + async run(input: LocalAgentRunInput, observer?: LocalAgentObserver): Promise { const runtime = await createCodexSdkLocalAgentRuntime(); - return runtime.run(input); + return runtime.run(input, observer); } } class ClaudeLocalAgentAdapter implements LocalAgentAdapter { readonly provider = "claude" as const; - async run(input: LocalAgentRunInput): Promise { + async run(input: LocalAgentRunInput, observer?: LocalAgentObserver): Promise { const { query } = await import("@anthropic-ai/claude-agent-sdk"); const claudeExecutable = process.env.CLAUDE_COMMAND ?? resolveExecutable("claude"); try { @@ -107,7 +111,11 @@ class ClaudeLocalAgentAdapter implements LocalAgentAdapter { for await (const message of messages) { items.push(message); const record = message as Record; - if (typeof record.session_id === "string") providerSessionId = record.session_id; + if (typeof record.session_id === "string") { + providerSessionId = record.session_id; + observer?.onSession?.(record.session_id); + } + notifyClaudeActivity(record, observer); if (record.type !== "result") continue; const resultError = claudeResultError(record); if (resultError) throw new Error(resultError); @@ -116,14 +124,21 @@ class ClaudeLocalAgentAdapter implements LocalAgentAdapter { finalResponse = extracted.finalResponse; structured = extracted.structured; } + const usage = claudeUsage(record.usage, "final"); + if (usage) observer?.onUsage?.(usage); } finalResponse = requireFinalResponse("Claude", finalResponse); + const usage = [...items] + .reverse() + .map((item) => claudeUsage((item as Record).usage, "final")) + .find((snapshot) => snapshot !== undefined); return { provider: this.provider, providerSessionId, finalResponse, items, + ...(usage ? { usage } : {}), ...(structured !== undefined ? { structured } : {}), }; } catch (error) { @@ -135,6 +150,58 @@ class ClaudeLocalAgentAdapter implements LocalAgentAdapter { } } +function notifyClaudeActivity(record: Record, observer?: LocalAgentObserver): void { + if (record.type === "tool_progress" && typeof record.tool_name === "string") { + observer?.onActivity?.({ kind: "tool", status: "running", label: record.tool_name }); + return; + } + if (record.type === "tool_use_summary" && typeof record.summary === "string") { + observer?.onActivity?.({ kind: "tool", status: "completed", label: record.summary }); + return; + } + if (record.type !== "assistant") return; + const message = record.message as { content?: unknown[] } | undefined; + for (const block of message?.content ?? []) { + const content = block as Record; + if (content.type !== "tool_use" || typeof content.name !== "string") continue; + observer?.onActivity?.({ + kind: content.name === "Bash" ? "command" : content.name === "Write" || content.name === "Edit" ? "file" : "tool", + status: "running", + label: content.name, + detail: claudeToolDetail(content.input), + }); + } +} + +function claudeToolDetail(input: unknown): string | undefined { + if (!input || typeof input !== "object") return undefined; + const record = input as Record; + for (const key of ["command", "file_path", "path", "query"]) { + if (typeof record[key] === "string") return record[key]; + } + return undefined; +} + +function claudeUsage(value: unknown, state: "partial" | "final"): LocalAgentUsageSnapshot | undefined { + if (!value || typeof value !== "object") return undefined; + const usage = value as Record; + const inputTokens = nonNegativeInteger(usage.input_tokens); + const outputTokens = nonNegativeInteger(usage.output_tokens); + if (inputTokens === undefined && outputTokens === undefined) return undefined; + return { + inputTokens, + cachedInputTokens: nonNegativeInteger(usage.cache_read_input_tokens), + cacheCreationInputTokens: nonNegativeInteger(usage.cache_creation_input_tokens), + outputTokens, + totalTokens: (inputTokens ?? 0) + (outputTokens ?? 0), + state, + }; +} + +function nonNegativeInteger(value: unknown): number | undefined { + return typeof value === "number" && Number.isSafeInteger(value) && value >= 0 ? value : undefined; +} + /** Build Claude SDK outputFormat when a JSON Schema is requested. */ export function claudeOutputFormatOptions( schema: JsonSchema | undefined, diff --git a/src/local-agent-runtime.test.ts b/src/local-agent-runtime.test.ts index 29b3d69f..055034a8 100644 --- a/src/local-agent-runtime.test.ts +++ b/src/local-agent-runtime.test.ts @@ -1,5 +1,5 @@ import assert from "node:assert/strict"; -import type { RunResult, ThreadOptions } from "@openai/codex-sdk"; +import type { RunResult, RunStreamedResult, ThreadEvent, ThreadOptions } from "@openai/codex-sdk"; import { CodexSdkLocalAgentRuntime, createCodexSdkLocalAgentRuntime, @@ -69,6 +69,63 @@ assert.deepEqual(codex.started[0], { modelReasoningEffort: undefined, }); +const streamedEvents: ThreadEvent[] = [ + { type: "thread.started", thread_id: "stream-thread" }, + { + type: "item.started", + item: { + id: "command-1", + type: "command_execution", + command: "npm test", + aggregated_output: "", + status: "in_progress", + }, + }, + { + type: "item.completed", + item: { + id: "command-1", + type: "command_execution", + command: "npm test", + aggregated_output: "ok", + exit_code: 0, + status: "completed", + }, + }, + { type: "item.completed", item: { id: "message-1", type: "agent_message", text: "done" } }, + { + type: "turn.completed", + usage: { input_tokens: 100, cached_input_tokens: 20, output_tokens: 30, reasoning_output_tokens: 10 }, + }, +]; +const streamingThread = { + id: "stream-thread", + async run(): Promise { throw new Error("unreachable"); }, + async runStreamed(): Promise { + return { events: (async function* () { yield* streamedEvents; })() }; + }, +}; +const observedSessions: string[] = []; +const observedActivity: string[] = []; +const observedUsage: number[] = []; +const streamedRuntime = new CodexSdkLocalAgentRuntime({ + startThread: () => streamingThread, + resumeThread: () => streamingThread, +}); +const streamed = await streamedRuntime.run( + { prompt: "test", workspace: "/tmp/project" }, + { + onSession: (id) => observedSessions.push(id), + onActivity: (activity) => observedActivity.push(`${activity.status}:${activity.label}`), + onUsage: (usage) => observedUsage.push(usage.totalTokens), + }, +); +assert.equal(streamed.finalResponse, "done"); +assert.equal(streamed.usage?.totalTokens, 130); +assert.deepEqual(observedSessions, ["stream-thread"]); +assert.deepEqual(observedActivity, ["running:npm test", "completed:npm test"]); +assert.deepEqual(observedUsage, [130]); + await runtime.run({ prompt: "make change", workspace: "/tmp/project", diff --git a/src/local-agent-runtime.ts b/src/local-agent-runtime.ts index e43a4aff..b3d7c69b 100644 --- a/src/local-agent-runtime.ts +++ b/src/local-agent-runtime.ts @@ -3,7 +3,10 @@ import type { CodexOptions, ModelReasoningEffort, RunResult, + RunStreamedResult, SandboxMode, + ThreadEvent, + ThreadItem, ThreadOptions, TurnOptions, } from "@openai/codex-sdk"; @@ -41,16 +44,40 @@ export interface LocalAgentRunResult { items: unknown[]; /** Provider-native structured object when schema was requested. */ structured?: unknown; + usage?: LocalAgentUsageSnapshot; +} + +export interface LocalAgentUsageSnapshot { + inputTokens?: number; + cachedInputTokens?: number; + cacheCreationInputTokens?: number; + outputTokens?: number; + totalTokens: number; + state: "partial" | "final"; +} + +export interface LocalAgentActivity { + kind: "tool" | "command" | "file" | "status"; + status: "running" | "completed" | "failed"; + label: string; + detail?: string; +} + +export interface LocalAgentObserver { + onSession?(providerSessionId: string): void; + onUsage?(usage: LocalAgentUsageSnapshot): void; + onActivity?(activity: LocalAgentActivity): void; } export interface LocalAgentRuntime { readonly provider: LocalAgentProvider; - run(input: LocalAgentRunInput): Promise; + run(input: LocalAgentRunInput, observer?: LocalAgentObserver): Promise; } interface CodexThreadLike { readonly id: string | null; run(prompt: string, turnOptions?: TurnOptions): Promise; + runStreamed?(prompt: string, turnOptions?: TurnOptions): Promise; } interface CodexClientLike { @@ -90,15 +117,18 @@ export class CodexSdkLocalAgentRuntime implements LocalAgentRuntime { this.codex = codex; } - async run(input: LocalAgentRunInput): Promise { + async run(input: LocalAgentRunInput, observer?: LocalAgentObserver): Promise { const options = threadOptionsFor(input); const thread = input.providerSessionId ? this.codex.resumeThread(input.providerSessionId, options) : this.codex.startThread(options); const turnOptions = input.schema ? { outputSchema: input.schema } : undefined; let turn: RunResult; + const streamed = thread.runStreamed !== undefined; try { - turn = await thread.run(input.prompt, turnOptions); + turn = thread.runStreamed + ? await collectCodexStream(await thread.runStreamed(input.prompt, turnOptions), observer) + : await thread.run(input.prompt, turnOptions); } catch (error) { if (input.schema && isNativeSchemaUnsupportedFailure(error)) { throw new ProviderSchemaUnsupportedError(this.provider, error); @@ -106,16 +136,90 @@ export class CodexSdkLocalAgentRuntime implements LocalAgentRuntime { throw error; } + if (!streamed && thread.id) observer?.onSession?.(thread.id); + const usage = turn.usage ? codexUsage(turn.usage) : undefined; + if (usage) observer?.onUsage?.(usage); return { provider: this.provider, providerSessionId: thread.id, finalResponse: turn.finalResponse, items: turn.items, + usage, ...(input.schema ? { structured: tryParseJson(turn.finalResponse) } : {}), }; } } +async function collectCodexStream( + streamed: RunStreamedResult, + observer?: LocalAgentObserver, +): Promise { + const items: ThreadItem[] = []; + let finalResponse = ""; + let usage: RunResult["usage"] = null; + for await (const event of streamed.events) { + if (event.type === "thread.started") observer?.onSession?.(event.thread_id); + if (event.type === "item.started") notifyCodexItem(event.item, "running", observer); + if (event.type === "item.completed") { + items.push(event.item); + notifyCodexItem(event.item, codexItemStatus(event.item), observer); + if (event.item.type === "agent_message") finalResponse = event.item.text; + } + if (event.type === "turn.completed") usage = event.usage; + if (event.type === "turn.failed") throw new Error(event.error.message); + if (event.type === "error") throw new Error(event.message); + } + return { items, finalResponse, usage }; +} + +function notifyCodexItem( + item: ThreadItem, + status: LocalAgentActivity["status"], + observer?: LocalAgentObserver, +): void { + const activity = codexItemActivity(item, status); + if (activity) observer?.onActivity?.(activity); +} + +function codexItemActivity( + item: ThreadItem, + status: LocalAgentActivity["status"], +): LocalAgentActivity | undefined { + if (item.type === "command_execution") { + return { kind: "command", status, label: item.command }; + } + if (item.type === "file_change") { + return { + kind: "file", + status, + label: "apply file changes", + detail: item.changes.map((change) => `${change.kind} ${change.path}`).join(", "), + }; + } + if (item.type === "mcp_tool_call") { + return { kind: "tool", status, label: `${item.server}.${item.tool}` }; + } + if (item.type === "web_search") return { kind: "tool", status, label: "web search", detail: item.query }; + return undefined; +} + +function codexItemStatus(item: ThreadItem): LocalAgentActivity["status"] { + if (item.type === "command_execution" || item.type === "mcp_tool_call" || item.type === "file_change") { + return item.status === "failed" ? "failed" : "completed"; + } + return "completed"; +} + +function codexUsage(usage: NonNullable): LocalAgentUsageSnapshot { + return { + inputTokens: usage.input_tokens, + cachedInputTokens: usage.cached_input_tokens, + outputTokens: usage.output_tokens, + totalTokens: usage.input_tokens + usage.output_tokens, + state: "final", + }; +} + function tryParseJson(text: string): unknown | undefined { try { return JSON.parse(text) as unknown; diff --git a/src/workflow-agent-observer.ts b/src/workflow-agent-observer.ts new file mode 100644 index 00000000..5e24cb03 --- /dev/null +++ b/src/workflow-agent-observer.ts @@ -0,0 +1,69 @@ +import type { LocalAgentObserver, LocalAgentUsageSnapshot } from "./local-agent-runtime.js"; +import type { WorkflowStore } from "./workflow-store.js"; + +const USAGE_WRITE_INTERVAL_MS = 5_000; + +export function createWorkflowAgentObserver( + store: WorkflowStore, + runId: string, + callIndex: number, + intervalMs = USAGE_WRITE_INTERVAL_MS, +): LocalAgentObserver & { close(): void } { + const baseline = store.getAgentCall(runId, callIndex)?.usage; + let lastUsageWrite = 0; + let pendingUsage: LocalAgentUsageSnapshot | undefined; + let timer: NodeJS.Timeout | undefined; + + const persistUsage = (usage: LocalAgentUsageSnapshot): void => { + pendingUsage = undefined; + if (timer) clearTimeout(timer); + timer = undefined; + lastUsageWrite = Date.now(); + store.updateAgentUsage(runId, callIndex, { + inputTokens: sumOptional(baseline?.inputTokens, usage.inputTokens), + cachedInputTokens: sumOptional(baseline?.cachedInputTokens, usage.cachedInputTokens), + cacheCreationInputTokens: sumOptional( + baseline?.cacheCreationInputTokens, + usage.cacheCreationInputTokens, + ), + outputTokens: sumOptional(baseline?.outputTokens, usage.outputTokens), + totalTokens: (baseline?.totalTokens ?? 0) + usage.totalTokens, + state: usage.state, + }); + }; + + const scheduleUsage = (): void => { + if (timer) return; + const wait = Math.max(0, intervalMs - (Date.now() - lastUsageWrite)); + timer = setTimeout(() => { + if (pendingUsage) persistUsage(pendingUsage); + }, wait); + timer.unref(); + }; + + return { + onSession(providerSessionId) { + store.attachAgentSession(runId, callIndex, providerSessionId); + }, + onActivity(activity) { + store.appendAgentActivity({ runId, callIndex, ...activity }); + }, + onUsage(usage) { + if (usage.state === "final" || Date.now() - lastUsageWrite >= intervalMs) { + persistUsage(usage); + return; + } + pendingUsage = usage; + scheduleUsage(); + }, + close() { + if (pendingUsage) persistUsage(pendingUsage); + if (timer) clearTimeout(timer); + timer = undefined; + }, + }; +} + +function sumOptional(left: number | undefined, right: number | undefined): number | undefined { + return left === undefined && right === undefined ? undefined : (left ?? 0) + (right ?? 0); +} diff --git a/src/workflow-api.ts b/src/workflow-api.ts index d30f4df2..158bc2c8 100644 --- a/src/workflow-api.ts +++ b/src/workflow-api.ts @@ -34,6 +34,7 @@ export { WorkflowEngineError } from "./workflow-errors.js"; // --------------------------------------------------------------------------- export interface WorkflowProviderRunInput { + callIndex: number; provider: LocalAgentProvider; prompt: string; providerSessionId?: string; @@ -388,6 +389,7 @@ export function createWorkflowApi(deps: WorkflowApiDeps): WorkflowApi { const cwd = worktreePath ?? deps.workspaceRoot; const providerBase = { + callIndex: index, provider, prompt: providerPrompt, model, diff --git a/src/workflow-store.test.ts b/src/workflow-store.test.ts index 935662c6..30c237e1 100644 --- a/src/workflow-store.test.ts +++ b/src/workflow-store.test.ts @@ -4,6 +4,7 @@ import { mkdtempSync, rmSync } from "node:fs"; import { tmpdir } from "node:os"; import { join } from "node:path"; import { openDatabase } from "./db/client.js"; +import { createWorkflowAgentObserver } from "./workflow-agent-observer.js"; import { WorkflowStore } from "./workflow-store.js"; const root = mkdtempSync(join(tmpdir(), "devspace-workflow-store-test-")); @@ -404,6 +405,37 @@ try { assert.ok(store.listRuns().length >= 3); + const observedRun = store.createRun({ + name: "Observe provider", + source: "inline", + scriptPath: join(root, "observe.js"), + scriptHash: "observer-test", + workspaceRoot: join(root, "project"), + }); + store.startAgentCall({ + runId: observedRun.id, + callIndex: 0, + cacheKey: "observer", + prompt: "Inspect the project", + provider: "codex", + }); + const observer = createWorkflowAgentObserver(store, observedRun.id, 0, 60_000); + observer.onSession?.("session_123"); + observer.onActivity?.({ kind: "command", status: "running", label: "npm test" }); + observer.onUsage?.({ inputTokens: 100, outputTokens: 20, totalTokens: 120, state: "partial" }); + observer.onUsage?.({ inputTokens: 180, outputTokens: 40, totalTokens: 220, state: "final" }); + observer.close(); + + const retryObserver = createWorkflowAgentObserver(store, observedRun.id, 0, 60_000); + retryObserver.onUsage?.({ inputTokens: 50, outputTokens: 30, totalTokens: 80, state: "final" }); + retryObserver.close(); + assert.equal(store.getAgentCall(observedRun.id, 0)?.providerSessionId, "session_123"); + assert.equal(store.getAgentCall(observedRun.id, 0)?.usage?.inputTokens, 230); + assert.equal(store.getAgentCall(observedRun.id, 0)?.usage?.outputTokens, 70); + assert.equal(store.getAgentCall(observedRun.id, 0)?.usage?.totalTokens, 300); + assert.equal(store.getAgentCall(observedRun.id, 0)?.usage?.state, "final"); + assert.equal(store.listAgentActivity(observedRun.id, 0)[0]?.label, "npm test"); + // Second store instance sees same rows const other = new WorkflowStore(root); stores.push(other); diff --git a/src/workflow-worker.ts b/src/workflow-worker.ts index ec442508..f2ded0a4 100644 --- a/src/workflow-worker.ts +++ b/src/workflow-worker.ts @@ -4,6 +4,7 @@ import { availableParallelism } from "node:os"; import type { ServerConfig } from "./config.js"; import { parseJsonText, type JsonValue } from "./json-types.js"; import { runLocalAgentProviderResult } from "./local-agent-adapters.js"; +import { createWorkflowAgentObserver } from "./workflow-agent-observer.js"; import { isLocalAgentProvider, loadLocalAgentProfiles, @@ -98,15 +99,25 @@ export async function runWorkflowWorker( if (abort.signal.aborted || store.isCancelRequested(runId)) { throw Object.assign(new Error("Workflow cancelled"), { name: "AbortError" }); } - const providerRun = await runLocalAgentProviderResult(input.provider, { - prompt: input.prompt, - workspace: input.workspace, - providerSessionId: input.providerSessionId, - model: input.model, - effort: input.effort, - writeMode: "allowed", - schema: input.schema, - }); + const observer = createWorkflowAgentObserver(store, runId, input.callIndex); + let providerRun; + try { + providerRun = await runLocalAgentProviderResult( + input.provider, + { + prompt: input.prompt, + workspace: input.workspace, + providerSessionId: input.providerSessionId, + model: input.model, + effort: input.effort, + writeMode: "allowed", + schema: input.schema, + }, + observer, + ); + } finally { + observer.close(); + } if (providerRun.isErr()) throw providerRun.error; const providerResult = providerRun.value; return {