terramend 0.2.0

package/src/agents/sessionLabeler.test.ts

@@ -0,0 +1,247 @@
+import { describe, expect, test } from "vitest";
+import {
+  deriveLabelFromTaskInput,
+  formatWithLabel,
+  ORCHESTRATOR_LABEL,
+  SessionLabeler,
+} from "#app/agents/sessionLabeler";
+
+describe("deriveLabelFromTaskInput", () => {
+  test("prefers explicit lens marker in prompt over description", () => {
+    expect(
+      deriveLabelFromTaskInput({
+        prompt: "lens: security\nReview the diff for...",
+        description: "general review",
+      }),
+    ).toBe("lens:security");
+  });
+
+  test("supports lens=<name> alternative syntax", () => {
+    expect(
+      deriveLabelFromTaskInput({
+        prompt: "lens=user-journey\nWalk through the happy path...",
+      }),
+    ).toBe("lens:user-journey");
+  });
+
+  test("falls back to description when no lens marker present", () => {
+    expect(
+      deriveLabelFromTaskInput({
+        prompt: "Review this diff for any bugs",
+        description: "Auth lens",
+      }),
+    ).toBe("lens:auth-lens");
+  });
+
+  test("falls back to subagent_type when description and lens marker absent", () => {
+    expect(
+      deriveLabelFromTaskInput({
+        prompt: "Some generic prompt",
+        subagent_type: "review",
+      }),
+    ).toBe("review");
+  });
+
+  test("returns generic subagent when nothing identifiable", () => {
+    expect(deriveLabelFromTaskInput({})).toBe("subagent");
+  });
+
+  test("slug normalizes whitespace and special chars", () => {
+    expect(
+      deriveLabelFromTaskInput({
+        description: "Schema migration & operational readiness!",
+      }),
+    ).toBe("lens:schema-migration-operational-readiness");
+  });
+
+  test("slug truncates labels longer than 40 chars to keep prefix readable", () => {
+    expect(
+      deriveLabelFromTaskInput({
+        description: "this is a very long lens description that exceeds the slug limit",
+      }),
+    ).toBe("lens:this-is-a-very-long-lens-description-tha");
+  });
+
+  test("ignores lens marker mid-line — must be at line start", () => {
+    expect(
+      deriveLabelFromTaskInput({
+        prompt: "Please review the lens: security claim made above",
+        description: "billing",
+      }),
+    ).toBe("lens:billing");
+  });
+});
+
+describe("SessionLabeler", () => {
+  test("first session seen is the orchestrator", () => {
+    const labeler = new SessionLabeler();
+    expect(labeler.labelFor("ses-A")).toBe(ORCHESTRATOR_LABEL);
+    // bound — same session returns same label on second call
+    expect(labeler.labelFor("ses-A")).toBe(ORCHESTRATOR_LABEL);
+    expect(labeler.size()).toBe(1);
+  });
+
+  test("FIFO matches dispatched labels to new sessions in dispatch order", () => {
+    const labeler = new SessionLabeler();
+    // orchestrator session
+    labeler.labelFor("parent");
+
+    // orchestrator dispatches 3 tasks in one assistant turn
+    labeler.recordTaskDispatch({ description: "security" });
+    labeler.recordTaskDispatch({ description: "correctness" });
+    labeler.recordTaskDispatch({ description: "user journey" });
+
+    expect(labeler.pendingDispatchCount()).toBe(3);
+
+    // children appear (potentially interleaved)
+    expect(labeler.labelFor("child-1")).toBe("lens:security");
+    expect(labeler.labelFor("child-2")).toBe("lens:correctness");
+    expect(labeler.labelFor("child-3")).toBe("lens:user-journey");
+
+    expect(labeler.pendingDispatchCount()).toBe(0);
+    expect(labeler.size()).toBe(4);
+  });
+
+  test("interleaved events from parent and children resolve to stable labels", () => {
+    const labeler = new SessionLabeler();
+    labeler.labelFor("parent");
+    labeler.recordTaskDispatch({ description: "security" });
+    labeler.recordTaskDispatch({ description: "correctness" });
+
+    // child-1 emits an event first (its label binds)
+    expect(labeler.labelFor("child-1")).toBe("lens:security");
+    // parent emits some events in between
+    expect(labeler.labelFor("parent")).toBe(ORCHESTRATOR_LABEL);
+    // child-2 finally appears
+    expect(labeler.labelFor("child-2")).toBe("lens:correctness");
+    // child-1 emits more events — still the same label
+    expect(labeler.labelFor("child-1")).toBe("lens:security");
+  });
+
+  test("falls back to subagent#N when child appears without a queued dispatch", () => {
+    const labeler = new SessionLabeler();
+    labeler.labelFor("parent");
+    // no recordTaskDispatch — but a child appears anyway (defensive path)
+    expect(labeler.labelFor("ghost")).toBe("subagent#1");
+    expect(labeler.labelFor("ghost-2")).toBe("subagent#2");
+  });
+
+  test("undefined/null/empty sessionID resolves to orchestrator label without binding", () => {
+    const labeler = new SessionLabeler();
+    expect(labeler.labelFor(undefined)).toBe(ORCHESTRATOR_LABEL);
+    expect(labeler.labelFor(null)).toBe(ORCHESTRATOR_LABEL);
+    expect(labeler.labelFor("")).toBe(ORCHESTRATOR_LABEL);
+    // size stays zero — those calls didn't bind anything
+    expect(labeler.size()).toBe(0);
+  });
+
+  test("entries returns insertion-ordered (sessionID, label) pairs", () => {
+    const labeler = new SessionLabeler();
+    labeler.labelFor("parent");
+    labeler.recordTaskDispatch({ description: "security" });
+    labeler.labelFor("child-1");
+    expect(labeler.entries()).toEqual([
+      ["parent", ORCHESTRATOR_LABEL],
+      ["child-1", "lens:security"],
+    ]);
+  });
+
+  test("Claude path: parent_tool_use_id resolves directly without consuming FIFO", () => {
+    // Claude runs subagents inside the orchestrator's session — they share
+    // session_id — and stamps subagent messages with parent_tool_use_id.
+    // recording dispatch with the Agent tool_use id binds it directly so
+    // future events resolve regardless of session_id.
+    const labeler = new SessionLabeler();
+    expect(labeler.labelFor("shared-session", null)).toBe(ORCHESTRATOR_LABEL);
+
+    labeler.recordTaskDispatch({ description: "correctness" }, "toolu_01");
+    labeler.recordTaskDispatch({ description: "security" }, "toolu_02");
+
+    // subagent events come through with shared session_id but distinct
+    // parent_tool_use_id — direct mapping wins
+    expect(labeler.labelFor("shared-session", "toolu_01")).toBe("lens:correctness");
+    expect(labeler.labelFor("shared-session", "toolu_02")).toBe("lens:security");
+
+    // orchestrator events on the same session still resolve correctly
+    expect(labeler.labelFor("shared-session", null)).toBe(ORCHESTRATOR_LABEL);
+
+    // pendingLabels is unused on the Claude path — FIFO never consumed
+    expect(labeler.pendingDispatchCount()).toBe(2);
+    expect(labeler.size()).toBe(1);
+  });
+
+  test("Claude path: unknown parent_tool_use_id falls through to sessionID/FIFO logic", () => {
+    // defensive: if a subagent event arrives with a parent_tool_use_id we
+    // never recorded (e.g. orchestrator dispatched off-stream, or a tool we
+    // didn't track), the labeler shouldn't crash — it should fall through
+    // to the sessionID-keyed path.
+    const labeler = new SessionLabeler();
+    labeler.labelFor("shared", null);
+    expect(labeler.labelFor("shared", "unknown-tool-id")).toBe(ORCHESTRATOR_LABEL);
+  });
+
+  test("realistic four-lens parallel fan-out — interleaved tool_use stream", () => {
+    // simulates the event order we'd see when the orchestrator dispatches
+    // 4 lens subagents in a single assistant turn and they all start emitting
+    // tool_use events more or less concurrently.
+    const labeler = new SessionLabeler();
+
+    // 1. orchestrator's `init` event
+    expect(labeler.labelFor("p")).toBe(ORCHESTRATOR_LABEL);
+
+    // 2. orchestrator emits 4 task tool_use events back-to-back
+    labeler.recordTaskDispatch({ description: "correctness & invariants" });
+    labeler.recordTaskDispatch({ description: "security" });
+    labeler.recordTaskDispatch({ description: "user journey" });
+    labeler.recordTaskDispatch({ description: "schema migration" });
+
+    // 3. children emit in arbitrary interleaved order
+    const observed: Array<[string, string]> = [];
+    for (const session of ["c1", "c2", "p", "c3", "c1", "c4", "c2", "p"]) {
+      observed.push([session, labeler.labelFor(session)]);
+    }
+
+    expect(observed).toEqual([
+      ["c1", "lens:correctness-invariants"],
+      ["c2", "lens:security"],
+      ["p", ORCHESTRATOR_LABEL],
+      ["c3", "lens:user-journey"],
+      ["c1", "lens:correctness-invariants"],
+      ["c4", "lens:schema-migration"],
+      ["c2", "lens:security"],
+      ["p", ORCHESTRATOR_LABEL],
+    ]);
+
+    expect(labeler.size()).toBe(5);
+    expect(labeler.pendingDispatchCount()).toBe(0);
+  });
+});
+
+describe("formatWithLabel", () => {
+  test("prefixes a single-line message with magenta-wrapped label", () => {
+    const out = formatWithLabel("orchestrator", "hello world");
+    expect(out).toContain("[orchestrator]");
+    expect(out).toContain("hello world");
+    // ANSI magenta + reset markers around the bracketed label (escapes
+    // built via fromCharCode to satisfy biome's no-control-character-in-regex)
+    const ESC = String.fromCharCode(27);
+    expect(out).toMatch(new RegExp(`${ESC}\\[35m\\[orchestrator\\]${ESC}\\[0m hello world$`));
+  });
+
+  test("prefixes every line of a multi-line message", () => {
+    const out = formatWithLabel("lens:security", "line one\nline two\nline three");
+    const lines = out.split("\n");
+    expect(lines).toHaveLength(3);
+    for (const line of lines) {
+      expect(line).toContain("[lens:security]");
+    }
+    expect(lines[0]).toContain("line one");
+    expect(lines[1]).toContain("line two");
+    expect(lines[2]).toContain("line three");
+  });
+
+  test("handles empty input without throwing", () => {
+    const out = formatWithLabel("orchestrator", "");
+    expect(out).toContain("[orchestrator]");
+  });
+});

package/src/agents/sessionLabeler.ts

@@ -0,0 +1,178 @@
+/**
+ * Track per-session labels so log lines from parallel subagents can be
+ * differentiated. The orchestrator dispatches lens subagents
+ * via the Task tool; each subagent runs in its own opencode/claude Session
+ * with its own `sessionID` (or `session_id`) tag on the NDJSON event stream.
+ *
+ * Without per-session prefixing, parallel subagent tool_use / tool_result /
+ * text events appear as a single interleaved stream tagged with `[Terramend]`,
+ * making it impossible for a human reading the logs to attribute work to a
+ * specific lens.
+ *
+ * The labeler is deliberately runtime-agnostic — both opencode.ts and
+ * claude.ts feed it the same shape. The contract is FIFO: when the orchestrator
+ * dispatches N task tool_use blocks in a single assistant turn (the parallel
+ * fan-out the multi-lens prompt requires), the i-th new sessionID is assumed
+ * to belong to the i-th task dispatch. This is correct as long as parallel
+ * dispatches are emitted in source-order and the runtimes respect that order
+ * when assigning child sessions; we do not depend on it for correctness of
+ * the read-only contract — only for log readability.
+ */
+
+export interface TaskDispatchInput {
+  description?: string | undefined;
+  subagent_type?: string | undefined;
+  prompt?: string | undefined;
+}
+
+export const ORCHESTRATOR_LABEL = "orchestrator";
+
+const LENS_PROMPT_PATTERN = /^\s*(?:lens|Lens|LENS)\s*[:=]\s*([A-Za-z][\w &/.-]{0,60})/m;
+
+function slug(value: string): string {
+  return value
+    .trim()
+    .toLowerCase()
+    .replace(/[^\w-]+/g, "-")
+    .replace(/^-+|-+$/g, "")
+    .slice(0, 40);
+}
+
+/**
+ * Extract a human-readable label from a Task tool's input. Tries (in order):
+ *   1. explicit `lens: <name>` marker on a line in the prompt — preferred,
+ *      lets the orchestrator name the lens deterministically
+ *   2. the Task tool's `description` field — short, written by orchestrator
+ *      per call, usually enough
+ *   3. the `subagent_type` — falls back to the named
+ *      subagent identity when description is missing
+ *   4. generic "subagent" — last resort
+ */
+export function deriveLabelFromTaskInput(input: TaskDispatchInput): string {
+  if (typeof input.prompt === "string") {
+    const match = input.prompt.match(LENS_PROMPT_PATTERN);
+    if (match?.[1]) {
+      const slugged = slug(match[1]);
+      if (slugged) return `lens:${slugged}`;
+    }
+  }
+  if (input.description) {
+    const slugged = slug(input.description);
+    if (slugged) return `lens:${slugged}`;
+  }
+  if (input.subagent_type) {
+    return input.subagent_type;
+  }
+  return "subagent";
+}
+
+/**
+ * Stateful tracker mapping subagent activity back to human-readable labels.
+ *
+ * Two attribution channels are supported because the runtimes differ:
+ *
+ *   - **OpenCode** spawns each subagent as its own opencode `Session` with
+ *     a distinct `sessionID`. The harness records each Task dispatch into a
+ *     pending FIFO queue; the next previously-unseen sessionID consumes the
+ *     head of the queue and binds it to that label.
+ *
+ *   - **Claude Code** runs subagents inside the orchestrator's session — they
+ *     all share `session_id` — and instead stamps every subagent message with
+ *     `parent_tool_use_id` pointing at the Agent tool_use id that spawned them.
+ *     The harness binds each Agent tool_use id to its dispatched label up
+ *     front, then `labelFor` looks the label up directly when an event arrives
+ *     carrying that `parent_tool_use_id`.
+ *
+ * `labelFor(sessionID, parentToolUseId?)` accepts both: when
+ * `parentToolUseId` is set and known it short-circuits to the direct mapping;
+ * otherwise it falls through to the FIFO/sessionID path.
+ */
+export class SessionLabeler {
+  private readonly labels = new Map<string, string>();
+  private readonly labelsByToolUseId = new Map<string, string>();
+  private readonly pendingLabels: string[] = [];
+  private fallbackCounter = 0;
+
+  /**
+   * Record a Task/Agent tool dispatch.
+   *
+   * @param input  Task tool input — used to derive the lens label.
+   * @param toolUseId  Optional Agent tool_use id. When provided, future events
+   *                   carrying `parent_tool_use_id === toolUseId` resolve
+   *                   directly to this label without consuming the FIFO queue
+   *                   (Claude path). Always also pushed to the FIFO queue so
+   *                   the OpenCode path still works when toolUseId is absent.
+   */
+  recordTaskDispatch(input: TaskDispatchInput, toolUseId?: string | null): string {
+    const label = deriveLabelFromTaskInput(input);
+    this.pendingLabels.push(label);
+    if (toolUseId) this.labelsByToolUseId.set(toolUseId, label);
+    return label;
+  }
+
+  /**
+   * Return a label for the given event.
+   *
+   * @param sessionID         Session id from the event (OpenCode: per-session;
+   *                          Claude: shared across orchestrator + subagents).
+   * @param parentToolUseId   Claude's `parent_tool_use_id` — non-null on
+   *                          subagent messages. When set and known, takes
+   *                          priority over the FIFO/sessionID path.
+   */
+  labelFor(sessionID: string | undefined | null, parentToolUseId?: string | null): string {
+    // Claude path: subagent messages carry parent_tool_use_id pointing at
+    // the Agent tool_use that spawned them. resolve directly without
+    // touching the sessionID-keyed map (which is bound to the orchestrator
+    // for the shared session_id and would otherwise misattribute).
+    if (parentToolUseId) {
+      const direct = this.labelsByToolUseId.get(parentToolUseId);
+      if (direct) return direct;
+    }
+
+    if (!sessionID) return ORCHESTRATOR_LABEL;
+    const existing = this.labels.get(sessionID);
+    if (existing) return existing;
+
+    let label: string;
+    if (this.labels.size === 0) {
+      label = ORCHESTRATOR_LABEL;
+    } else if (this.pendingLabels.length > 0) {
+      label = this.pendingLabels.shift() as string;
+    } else {
+      this.fallbackCounter += 1;
+      label = `subagent#${this.fallbackCounter}`;
+    }
+    this.labels.set(sessionID, label);
+    return label;
+  }
+
+  /** number of distinct sessions seen so far (for diagnostics) */
+  size(): number {
+    return this.labels.size;
+  }
+
+  /** all (sessionID, label) pairs, oldest first */
+  entries(): Array<[string, string]> {
+    return Array.from(this.labels.entries());
+  }
+
+  /** how many pending labels are queued waiting to bind to a new session */
+  pendingDispatchCount(): number {
+    return this.pendingLabels.length;
+  }
+}
+
+/**
+ * Format a log message with a session label prefix in magenta. Mirrors the
+ * style of utils/log.ts:prefixLines() so per-session prefixes look the same
+ * as the dormant withLogPrefix-based ones.
+ */
+export function formatWithLabel(label: string, message: string): string {
+  const MAGENTA = "\x1b[35m";
+  const RESET = "\x1b[0m";
+  const colored = `${MAGENTA}[${label}]${RESET} `;
+  return message
+    .split("\n")
+    .map((line) => `${colored}${line}`)
+    .join("\n");
+}

package/src/agents/shared.test.ts

@@ -0,0 +1,76 @@
+import { describe, expect, it } from "vitest";
+import { type AgentUsage, mergeAgentUsage } from "#app/agents/shared";
+
+const entry = (overrides: Partial<AgentUsage>): AgentUsage => ({
+  agent: "terramend",
+  inputTokens: 0,
+  outputTokens: 0,
+  ...overrides,
+});
+
+describe("mergeAgentUsage", () => {
+  it("returns undefined when both sides are undefined", () => {
+    expect(mergeAgentUsage(undefined, undefined)).toBeUndefined();
+  });
+
+  it("returns a copy of b when a is undefined", () => {
+    const b = entry({ inputTokens: 10 });
+    expect(mergeAgentUsage(undefined, b)).toEqual(b);
+  });
+
+  it("returns a copy of a when b is undefined", () => {
+    const a = entry({ inputTokens: 10 });
+    expect(mergeAgentUsage(a, undefined)).toEqual(a);
+  });
+
+  it("sums inputTokens and outputTokens unconditionally", () => {
+    const merged = mergeAgentUsage(
+      entry({ inputTokens: 10, outputTokens: 5 }),
+      entry({ inputTokens: 20, outputTokens: 7 }),
+    );
+    expect(merged?.inputTokens).toBe(30);
+    expect(merged?.outputTokens).toBe(12);
+  });
+
+  it("keeps cache/cost fields undefined when both sides lack them", () => {
+    // this matters so downstream aggregateUsage doesn't persist spurious 0s into the DB
+    const merged = mergeAgentUsage(entry({ inputTokens: 10 }), entry({ inputTokens: 20 }));
+    expect(merged?.cacheReadTokens).toBeUndefined();
+    expect(merged?.cacheWriteTokens).toBeUndefined();
+    expect(merged?.costUsd).toBeUndefined();
+  });
+
+  it("sums cache and cost fields when either side reports them", () => {
+    const merged = mergeAgentUsage(
+      entry({ inputTokens: 10, cacheReadTokens: 100, costUsd: 0.01 }),
+      entry({ inputTokens: 20, cacheWriteTokens: 50, costUsd: 0.02 }),
+    );
+    expect(merged?.cacheReadTokens).toBe(100);
+    expect(merged?.cacheWriteTokens).toBe(50);
+    expect(merged?.costUsd).toBeCloseTo(0.03, 10);
+  });
+
+  it("preserves the agent id of the left operand", () => {
+    // the aggregator is called inside a single agent's run() — the agent label
+    // is a fixed property of the harness, not something that can flip mid-run
+    const merged = mergeAgentUsage(
+      entry({ agent: "claude", inputTokens: 10 }),
+      entry({ agent: "something-else", inputTokens: 20 }),
+    );
+    expect(merged?.agent).toBe("claude");
+  });
+
+  it("returns a fresh object rather than the input reference", () => {
+    // callers treat AgentUsage as immutable; returning the input itself would
+    // leak that invariant. mutating the returned value must not affect inputs.
+    const a = entry({ inputTokens: 10 });
+    const mergedWithUndef = mergeAgentUsage(a, undefined);
+    expect(mergedWithUndef).not.toBe(a);
+    expect(mergedWithUndef).toEqual(a);
+
+    const b = entry({ inputTokens: 20 });
+    const mergedFromUndef = mergeAgentUsage(undefined, b);
+    expect(mergedFromUndef).not.toBe(b);
+    expect(mergedFromUndef).toEqual(b);
+  });
+});