macroclaw 0.8.0 → 0.9.0

package/src/orchestrator.ts

@@ -1,16 +1,22 @@
 import { z } from "zod/v4";
-import { Claude, type ClaudeDeferredResult, ClaudeParseError, ClaudeProcessError, type ClaudeResult, isDeferred } from "./claude";
-import { logPrompt, logResult } from "./history";
+import {
+  Claude,
+  QueryParseError,
+  QueryProcessError,
+  type QueryResult,
+  QueryValidationError,
+  type RunningQuery
+} from "./claude";
+import { writeHistoryPrompt, writeHistoryResult } from "./history";
 import { createLogger } from "./logger";
-import { BG_TIMEOUT, CRON_TIMEOUT, MAIN_TIMEOUT, SYSTEM_PROMPT } from "./prompts";
+import { CRON_TIMEOUT, MAIN_TIMEOUT, SYSTEM_PROMPT } from "./prompts";
 import { Queue } from "./queue";
-import { loadSessions, newSessionId, type Sessions, saveSessions } from "./sessions";
+import { loadSessions, saveSessions } from "./sessions";
 import type { ButtonSpec } from "./telegram";
 
 const log = createLogger("orchestrator");
 
-// --- Response schema (owned by orchestrator) ---
-// Flat object (no oneOf/discriminatedUnion) — Claude CLI --json-schema requires a single top-level object.
+// --- Response schema ---
 
 const backgroundAgentSchema = z.object({
   name: z.string().describe("Label for the background agent"),
@@ -18,7 +24,7 @@ const backgroundAgentSchema = z.object({
   model: z.enum(["haiku", "sonnet", "opus"]).describe("Model to use for the background agent").optional(),
 });
 
-const claudeResponseSchema = z.object({
+const agentOutputSchema = z.object({
   action: z.enum(["send", "silent"]).describe("'send' to reply to the user, 'silent' to do nothing"),
   actionReason: z.string().describe("Why the agent chose this action (logged, not sent)"),
   message: z.string().describe("The message to send to Telegram (required when action is 'send')").optional(),
@@ -27,9 +33,11 @@ const claudeResponseSchema = z.object({
   backgroundAgents: z.array(backgroundAgentSchema).describe("Background agents to spawn alongside this response").optional(),
 });
 
-type ClaudeResponseInternal = z.infer<typeof claudeResponseSchema>;
+type AgentOutput = z.infer<typeof agentOutputSchema>;
 
-const jsonSchema = JSON.stringify(z.toJSONSchema(claudeResponseSchema, { target: "jsonSchema7" }));
+const responseResultType = { type: "object" as const, schema: agentOutputSchema };
+
+const textResultType = { type: "text" } as const;
 
 // --- Public response type ---
 
@@ -46,8 +54,7 @@ export interface OrchestratorResponse {
 type OrchestratorRequest =
   | { type: "user"; message: string; files?: string[] }
   | { type: "cron"; name: string; prompt: string; model?: string }
-  | { type: "background-agent-result"; name: string; response: ClaudeResponseInternal; sessionId?: string }
-  | { type: "background-agent"; name: string; prompt: string; model?: string }
+  | { type: "background-agent-result"; name: string; response: AgentOutput; sessionId?: string }
   | { type: "button"; label: string };
 
 function escapeHtml(text: string): string {
@@ -62,8 +69,6 @@ interface BackgroundInfo {
   startTime: Date;
 }
 
-// --- Orchestrator ---
-
 export interface OrchestratorConfig {
   model?: string;
   workspace: string;
@@ -72,45 +77,21 @@ export interface OrchestratorConfig {
   claude?: Claude;
 }
 
-interface BuiltRequest {
-  prompt: string;
-  model: string | undefined;
-  timeout: number;
-  files?: string[];
-  useMainSession: boolean;
-}
-
-interface CallResult {
-  response: ClaudeResponseInternal;
-  sessionId: string;
-}
-
 export class Orchestrator {
+  #config: Omit<OrchestratorConfig , 'claude'>;
   #claude: Claude;
-  #sessions: Sessions;
-  #sessionId: string;
-  #resumeSession: boolean;
-  #sessionResolved = false;
-  #config: OrchestratorConfig;
-  #active = new Map<string, BackgroundInfo>();
+
+  #mainSessionId: string | undefined;
+  #backgroundAgents = new Map<string, BackgroundInfo>();
   #queue: Queue<OrchestratorRequest>;
 
   constructor(config: OrchestratorConfig) {
     this.#config = config;
-    this.#claude = config.claude ?? new Claude({ workspace: config.workspace, jsonSchema });
-    this.#sessions = loadSessions(config.settingsDir);
+    this.#claude = config.claude ?? new Claude({ workspace: config.workspace, systemPrompt: SYSTEM_PROMPT });
     this.#queue = new Queue<OrchestratorRequest>();
     this.#queue.setHandler((request) => this.#handleRequest(request));
 
-    if (this.#sessions.mainSessionId) {
-      this.#sessionId = this.#sessions.mainSessionId;
-      this.#resumeSession = true;
-    } else {
-      this.#sessionId = newSessionId();
-      this.#resumeSession = false;
-      saveSessions({ mainSessionId: this.#sessionId }, config.settingsDir);
-      log.info({ sessionId: this.#sessionId }, "Created new session");
-    }
+    this.#mainSessionId = loadSessions(config.settingsDir).mainSessionId;
   }
 
   // --- Public handle methods ---
@@ -134,7 +115,7 @@ export class Orchestrator {
   }
 
   handleBackgroundList(): void {
-    const agents = [...this.#active.values()];
+    const agents = [...this.#backgroundAgents.values()];
     if (agents.length === 0) {
       this.#callOnResponse({ message: "No background agents running." });
       return;
@@ -148,12 +129,12 @@ export class Orchestrator {
       const text = `${a.name} (${elapsed}s)`.slice(0, 27);
       return { text, data: `peek:${a.sessionId}` };
     });
-    buttons.push("_dismiss");
+    buttons.push({text: "Dismiss", data: "_dismiss"});
     this.#callOnResponse({ message: lines.join("\n"), buttons });
   }
 
   async handlePeek(sessionId: string): Promise<void> {
-    const agent = this.#active.get(sessionId);
+    const agent = this.#backgroundAgents.get(sessionId);
     if (!agent) {
       this.#callOnResponse({ message: "Agent not found or already finished." });
       return;
@@ -162,24 +143,21 @@ export class Orchestrator {
     this.#callOnResponse({ message: `Peeking at <b>${escapeHtml(agent.name)}</b>...` });
 
     try {
-      const result = await this.#claude.run({
-        resume: true,
+      const query = this.#claude.forkSession(
         sessionId,
-        forkSession: true,
-        model: "haiku",
-        plainText: true,
-        prompt: "Give a brief status update: what has been done so far, what's currently happening, and what's remaining. 2-3 sentences max.",
-      });
-
-      const text = isDeferred(result) ? "Agent is still working..." : (result.result ?? "[No output]");
-      this.#callOnResponse({ message: `<b>[${escapeHtml(agent.name)}]</b> ${text}` });
+        "Give a brief status update: what has been done so far, what's currently happening, and what's remaining. 2-3 sentences max.",
+        textResultType,
+        { model: "haiku" },
+      );
+      const { value } = await query.result;
+      this.#callOnResponse({ message: `<b>[${escapeHtml(agent.name)}]</b> ${value || "[No output]"}` });
     } catch (err) {
       this.#callOnResponse({ message: `Couldn't peek at ${escapeHtml(agent.name)}: ${err}` });
     }
   }
 
   handleSessionCommand(): void {
-    this.#callOnResponse({ message: `Session: <code>${this.#sessionId}</code>` });
+    this.#callOnResponse({ message: `Session: <code>${this.#mainSessionId ?? "none"}</code>` });
   }
 
   // --- Internal queue handler ---
@@ -187,39 +165,32 @@ export class Orchestrator {
   async #handleRequest(request: OrchestratorRequest): Promise<void> {
     log.debug({ type: request.type }, "Incoming request");
 
-    // Background result with matching session ID: deliver directly without Claude round-trip
-    if (request.type === "background-agent-result" && "sessionId" in request && request.sessionId === this.#sessionId) {
+    // Background result with matching session ID: deliver directly
+    if (request.type === "background-agent-result" && request.sessionId === this.#mainSessionId) {
       log.debug({ name: request.name }, "Background result on current session, applying directly");
-      await this.#deliverClaudeResponse(request.response);
+      await this.#deliverResponse(request.response);
       return;
     }
 
-    // Fork session if a backgrounded task is running on the main session
-    const needsFork = (request.type === "user" || request.type === "button") && this.#active.has(this.#sessionId);
-
-    const startTime = new Date();
-    const rawResponse = await this.#processRequest(request, needsFork ? { forkSession: true } : undefined);
-    if (isDeferred(rawResponse)) {
-      const name = request.type === "user" ? request.message.slice(0, 30).replace(/\s+/g, "-")
-        : request.type === "cron" ? `cron-${request.name}`
-        : "task";
-      log.info({ name, sessionId: rawResponse.sessionId }, "Request backgrounded due to timeout");
-      this.#callOnResponse({ message: "This is taking longer, continuing in the background." });
-      this.#adoptBackground(name, rawResponse.sessionId, startTime, rawResponse.completion.then(
-        (r) => {
-          const msg = r.structuredOutput ? String((r.structuredOutput as Record<string, unknown>).message ?? "") : (r.result ?? "");
-          return { action: "send" as const, message: msg, actionReason: "deferred-completed" };
-        },
-        (err) => ({ action: "send" as const, message: `[Error] ${err}`, actionReason: "deferred-failed" }),
-      ));
-      return;
+    await writeHistoryPrompt(request);
+
+    const result = await this.#queryWithRetry(request);
+    if (!result) return;
+
+    // Update session ID (important after fork or new session)
+    if (result.sessionId !== this.#mainSessionId) {
+      log.info({ oldSessionId: this.#mainSessionId, newSessionId: result.sessionId }, "Session updated");
+      this.#mainSessionId = result.sessionId;
+      saveSessions({ mainSessionId: this.#mainSessionId }, this.#config.settingsDir);
     }
 
-    log.debug({ action: rawResponse.action, actionReason: rawResponse.actionReason }, "Response");
-    await this.#deliverClaudeResponse(rawResponse);
+    await writeHistoryResult(result.value);
+    await this.#deliverResponse(result.value);
   }
 
-  async #deliverClaudeResponse(response: ClaudeResponseInternal): Promise<void> {
+  // --- Response delivery ---
+
+  async #deliverResponse(response: AgentOutput): Promise<void> {
     if (response.action === "send") {
       this.#callOnResponse({
         message: response.message || "[No output]",
@@ -245,187 +216,151 @@ export class Orchestrator {
     });
   }
 
-  // --- Internal background management ---
+  // --- Claude calls ---
 
-  #spawnBackground(name: string, prompt: string, model: string | undefined) {
-    const sessionId = newSessionId();
-    const info: BackgroundInfo = { name, sessionId, startTime: new Date() };
-    this.#active.set(sessionId, info);
-
-    log.debug({ name, sessionId }, "Starting background agent");
+  async #queryWithRetry(request: OrchestratorRequest): Promise<QueryResult<AgentOutput> | null> {
+    const timeout = request.type === "cron" ? CRON_TIMEOUT : MAIN_TIMEOUT;
+    const query = this.#query(request);
 
-    this.#processRequest({ type: "background-agent", name, prompt, model }).then(
-      async (rawResponse) => {
-        let response: ClaudeResponseInternal;
-        if (isDeferred(rawResponse)) {
-          try {
-            const r = await rawResponse.completion;
-            response = { action: "send", message: String(r.structuredOutput ?? r.result ?? ""), actionReason: "deferred-completed" };
-          } catch (err) {
-            response = { action: "send", message: `[Error] ${err}`, actionReason: "deferred-failed" };
-          }
-        } else {
-          response = rawResponse;
+    try {
+      const result = await this.#awaitOrBackground(query, request, timeout);
+      if (!result) return null;
+      return result;
+    } catch (err) {
+      // Resume failed — retry with a fresh session
+      if (err instanceof QueryProcessError && this.#mainSessionId) {
+        log.info("Resume failed, retrying with new session");
+        this.#mainSessionId = undefined;
+        const retryQuery = this.#query(request);
+
+        try {
+          const retryResult = await this.#awaitOrBackground(retryQuery, request, timeout);
+          if (!retryResult) return null;
+          return retryResult;
+        } catch (retryErr) {
+          return { value: this.#errorResponse(retryErr), sessionId: retryQuery.sessionId };
         }
-        this.#active.delete(sessionId);
-        log.debug({ name, message: response.message }, "Background agent finished");
-        this.#queue.push({ type: "background-agent-result", name, response });
-      },
-      (err) => {
-        this.#active.delete(sessionId);
-        log.error({ name, err }, "Background agent failed");
-        this.#queue.push({ type: "background-agent-result", name, response: { action: "send", message: `[Error] ${err}`, actionReason: "bg-agent-failed" } });
-      },
-    );
-  }
-
-  #adoptBackground(name: string, sessionId: string, startTime: Date, completion: Promise<ClaudeResponseInternal>) {
-    const info: BackgroundInfo = { name, sessionId, startTime };
-    this.#active.set(sessionId, info);
-
-    log.debug({ name, sessionId }, "Adopting backgrounded task");
-
-    // completion is pre-wrapped with .then(ok, err) at the call site, so it never rejects
-    completion.then((response) => {
-      this.#active.delete(sessionId);
-      log.debug({ name, message: response.message }, "Adopted task finished");
-      this.#queue.push({ type: "background-agent-result", name, response, sessionId });
-    });
-  }
-
-  // --- Core Claude processing ---
-
-  async #processRequest(request: OrchestratorRequest, options?: { forkSession?: boolean }): Promise<ClaudeResponseInternal | ClaudeDeferredResult> {
-    const built = this.#buildRequest(request);
-    await logPrompt(request);
-
-    if (built.useMainSession) {
-      let result = await this.#callClaude(built, this.#resumeSession, this.#sessionId, options?.forkSession);
-
-      // Session resolution: if resume failed on first call, create new session
-      if (!isDeferred(result) && !this.#sessionResolved && this.#resumeSession && result.response.actionReason === "process-error") {
-        this.#sessionId = newSessionId();
-        log.info({ sessionId: this.#sessionId }, "Resume failed, created new session");
-        this.#resumeSession = false;
-        saveSessions({ mainSessionId: this.#sessionId }, this.#config.settingsDir);
-        result = await this.#callClaude(built, this.#resumeSession, this.#sessionId);
       }
 
-      if (isDeferred(result)) return result;
+      return { value: this.#errorResponse(err), sessionId: this.#mainSessionId ?? "" };
+    }
+  }
 
-      // Update session ID from response (important after fork)
-      if (result.sessionId && result.sessionId !== this.#sessionId) {
-        log.info({ oldSessionId: this.#sessionId, newSessionId: result.sessionId }, "Session forked, updating session ID");
-        this.#sessionId = result.sessionId;
-        saveSessions({ mainSessionId: this.#sessionId }, this.#config.settingsDir);
-      }
+  #query(request: OrchestratorRequest) {
+    const prompt = this.#formatPrompt(request);
+    const model = request.type === "cron" ? (request.model ?? this.#config.model) : this.#config.model;
+    const opts = { model };
 
-      // Mark resolved on first success
-      if (!this.#sessionResolved && result.response.actionReason !== "process-error") {
-        this.#sessionResolved = true;
-        this.#resumeSession = true;
-      }
+    // Fork if a background agent is running on the main session
+    if (this.#mainSessionId && this.#backgroundAgents.has(this.#mainSessionId)) {
+      return this.#claude.forkSession(this.#mainSessionId, prompt, responseResultType, opts);
+    }
 
-      await logResult(result.response);
-      return result.response;
+    // Resume existing session
+    if (this.#mainSessionId) {
+      return this.#claude.resumeSession(this.#mainSessionId, prompt, responseResultType, opts);
     }
 
-    // background-agent: fork from main session for full context
-    log.debug({ name: (request as { name: string }).name }, "Processing background-agent (forked session)");
-    const bgResult = await this.#callClaude(built, true, this.#sessionId, true);
-    if (isDeferred(bgResult)) return bgResult;
-    await logResult(bgResult.response);
-    return bgResult.response;
+    // Start fresh
+    return this.#claude.newSession(prompt, responseResultType, opts);
   }
 
-  #buildRequest(request: OrchestratorRequest): BuiltRequest {
+  #formatPrompt(request: OrchestratorRequest): string {
     switch (request.type) {
-      case "user":
-        return {
-          prompt: this.#buildPromptWithFiles(request.message, request.files),
-          model: this.#config.model,
-          timeout: MAIN_TIMEOUT,
-          useMainSession: true,
-        };
+      case "user": {
+        if (!request.files?.length) return request.message;
+        const prefix = request.files.map((f) => `[File: ${f}]`).join("\n");
+        return request.message ? `${prefix}\n${request.message}` : prefix;
+      }
       case "cron":
-        return {
-          prompt: `[Context: cron/${request.name}] ${request.prompt}`,
-          model: request.model ?? this.#config.model,
-          timeout: CRON_TIMEOUT,
-          useMainSession: true,
-        };
+        return `[Context: cron/${request.name}] ${request.prompt}`;
       case "background-agent-result":
-        return {
-          prompt: `[Context: background-result/${request.name}] ${request.response.message || "[No output]"}`,
-          model: this.#config.model,
-          timeout: MAIN_TIMEOUT,
-          useMainSession: true,
-        };
+        return `[Context: background-result/${request.name}] ${request.response.message || "[No output]"}`;
       case "button":
-        return {
-          prompt: `[Context: button-click] User tapped "${request.label}"`,
-          model: this.#config.model,
-          timeout: MAIN_TIMEOUT,
-          useMainSession: true,
-        };
-      case "background-agent":
-        return {
-          prompt: `[Context: background-agent/${request.name}] ${request.prompt}`,
-          model: request.model ?? this.#config.model,
-          timeout: BG_TIMEOUT,
-          useMainSession: false,
-        };
+        return `[Context: button-click] User tapped "${request.label}"`;
     }
   }
 
-  #buildPromptWithFiles(message: string, files?: string[]): string {
-    if (!files?.length) return message;
-    const prefix = files.map((f) => `[File: ${f}]`).join("\n");
-    return message ? `${prefix}\n${message}` : prefix;
+  async #awaitOrBackground(
+    query: RunningQuery<AgentOutput>,
+    request: OrchestratorRequest,
+    timeoutMs: number,
+  ): Promise<QueryResult<AgentOutput> | null> {
+    const result = await Promise.race([
+      query.result,
+      new Promise<null>((resolve) => setTimeout(() => resolve(null), timeoutMs)),
+    ]);
+
+    if (result !== null) return result;
+
+    const name = request.type === "user" ? request.message.slice(0, 30).replace(/\s+/g, "-")
+      : request.type === "cron" ? `cron-${request.name}`
+      : "task";
+    log.info({ name, sessionId: query.sessionId }, "Request backgrounded due to timeout");
+    this.#callOnResponse({ message: "This is taking longer, continuing in the background." });
+    this.#adoptBackground(name, query.sessionId, query.startedAt, query.result);
+    return null;
   }
 
-  #validateResponse(raw: unknown): ClaudeResponseInternal {
-    const parsed = claudeResponseSchema.safeParse(raw);
-    if (parsed.success) return parsed.data;
-    log.warn({ error: parsed.error.message }, "structured_output failed validation");
-    const rawObj = raw as Record<string, unknown> | null;
-    const msg = typeof rawObj?.message === "string" ? rawObj.message : JSON.stringify(raw);
-    return { action: "send", message: msg, actionReason: "validation-failed" };
+  #errorResponse(err: unknown): AgentOutput {
+    if (err instanceof QueryProcessError) {
+      return { action: "send", message: `[Error] Claude exited with code ${err.exitCode}:\n${err.stderr}`, actionReason: "process-error" };
+    }
+    if (err instanceof QueryParseError) {
+      return { action: "send", message: `[JSON Error] ${err.raw}`, actionReason: "json-parse-failed" };
+    }
+    if (err instanceof QueryValidationError) {
+      const rawObj = err.raw as Record<string, unknown> | null;
+      const msg = typeof rawObj?.message === "string" ? rawObj.message : JSON.stringify(err.raw);
+      return { action: "send", message: msg, actionReason: "validation-failed" };
+    }
+    throw err;
   }
 
-  #resultToCallResult(result: ClaudeResult): CallResult {
-    if (result.structuredOutput) {
-      return { response: this.#validateResponse(result.structuredOutput), sessionId: result.sessionId };
-    }
-    log.error({ hasResult: !!result.result }, "No structured_output in response");
-    const raw = result.result ? escapeHtml(result.result) : "";
-    const msg = raw ? `[No structured output] ${raw}` : "[No output]";
-    return { response: { action: "send", message: msg, actionReason: "no-structured-output" }, sessionId: result.sessionId };
+  // --- Background management ---
+
+  #spawnBackground(name: string, prompt: string, model: string | undefined) {
+    const bgPrompt = `[Context: background-agent/${name}] ${prompt}`;
+    const query = this.#mainSessionId
+      ? this.#claude.forkSession(this.#mainSessionId, bgPrompt, responseResultType, { model })
+      : this.#claude.newSession(bgPrompt, responseResultType, { model });
+    const sessionId = query.sessionId;
+    const info: BackgroundInfo = { name, sessionId, startTime: query.startedAt };
+    this.#backgroundAgents.set(sessionId, info);
+
+    log.debug({ name, sessionId }, "Starting background agent");
+
+    query.result.then(
+      async ({ value: response }) => {
+        this.#backgroundAgents.delete(sessionId);
+        log.debug({ name, message: response.message }, "Background agent finished");
+        this.#queue.push({ type: "background-agent-result", name, response });
+      },
+      (err) => {
+        this.#backgroundAgents.delete(sessionId);
+        log.error({ name, err }, "Background agent failed");
+        this.#queue.push({ type: "background-agent-result", name, response: { action: "send", message: `[Error] ${err}`, actionReason: "bg-agent-failed" } });
+      },
+    );
   }
 
-  async #callClaude(built: BuiltRequest, resume: boolean, sid: string, forkSession?: boolean): Promise<CallResult | ClaudeDeferredResult> {
-    try {
-      const result = await this.#claude.run({
-        prompt: built.prompt,
-        resume,
-        sessionId: sid,
-        forkSession,
-        model: built.model,
-        systemPrompt: SYSTEM_PROMPT,
-        timeoutMs: built.timeout,
-      });
+  #adoptBackground(name: string, sessionId: string, startTime: Date, completion: Promise<QueryResult<AgentOutput>>) {
+    const info: BackgroundInfo = { name, sessionId, startTime };
+    this.#backgroundAgents.set(sessionId, info);
 
-      if (isDeferred(result)) return result;
+    log.debug({ name, sessionId }, "Adopting backgrounded task");
 
-      return this.#resultToCallResult(result);
-    } catch (err) {
-      if (err instanceof ClaudeProcessError) {
-        return { response: { action: "send", message: `[Error] Claude exited with code ${err.exitCode}:\n${err.stderr}`, actionReason: "process-error" }, sessionId: sid };
-      }
-      if (err instanceof ClaudeParseError) {
-        return { response: { action: "send", message: `[JSON Error] ${err.raw}`, actionReason: "json-parse-failed" }, sessionId: sid };
-      }
-      throw err;
-    }
+    completion.then(
+      ({ value: response }) => {
+        this.#backgroundAgents.delete(sessionId);
+        log.debug({ name }, "Adopted task finished");
+        this.#queue.push({ type: "background-agent-result", name, response, sessionId });
+      },
+      (err) => {
+        this.#backgroundAgents.delete(sessionId);
+        log.error({ name, err }, "Adopted task failed");
+        this.#queue.push({ type: "background-agent-result", name, response: { action: "send", message: `[Error] ${err}`, actionReason: "deferred-failed" }, sessionId });
+      },
+    );
   }
 }