@saadnvd1/agent-os 0.3.471 → 0.3.472

package/lib/agents/launch.ts

@@ -4,6 +4,7 @@ import path from "path";
 import { execFile } from "child_process";
 import { promisify } from "util";
 import { getProvider } from "../providers";
+import { sessionShellEnv } from "../models/launch";
 import { shellQuote } from "../hosts/ssh";
 import { runInBackground } from "../async-operations";
 import { trustPromptKeys } from "../tasks/state";
@@ -64,6 +65,8 @@ export async function launchClaude(opts: {
   tmuxName: string;
   cwd: string;
   model: string;
+  // The model's provider (lib/models); Anthropic when absent.
+  provider?: string | null;
   prompt: string;
   brief?: string;
   // Continue this Claude conversation; the prompt is its next message.
@@ -88,7 +91,13 @@ export async function launchClaude(opts: {
       sessionId: opts.resume,
     })
     .join(" ");
-  const agent = `export PATH=${shellQuote(AOS_BIN_DIR)}:"$HOME/.local/bin:$PATH"; ${provider.command} ${flags} ${CLAUDE_STATUS_SETTINGS_FLAG} --append-system-prompt-file ${shellQuote(briefFile)} "$(cat ${shellQuote(promptFile)})"; exec "\${SHELL:-/bin/sh}" -l`;
+  // The provider's key is fetched by the pane's own shell from its
+  // reference: only the reference is in this command line.
+  const routing = sessionShellEnv({
+    model: opts.model,
+    model_provider: opts.provider ?? null,
+  });
+  const agent = `export PATH=${shellQuote(AOS_BIN_DIR)}:"$HOME/.local/bin:$PATH"; ${routing.exportKey} ${provider.command} ${flags} ${CLAUDE_STATUS_SETTINGS_FLAG} --append-system-prompt-file ${shellQuote(briefFile)} "$(cat ${shellQuote(promptFile)})"; exec "\${SHELL:-/bin/sh}" -l`;
   await execFileAsync(
     "tmux",
     [
@@ -98,7 +107,7 @@ export async function launchClaude(opts: {
       opts.tmuxName,
       "-c",
       opts.cwd,
-      ...envArgs(agentEnv(opts.sessionId)),
+      ...envArgs({ ...agentEnv(opts.sessionId), ...routing.env }),
       agent,
     ],
     { cwd: opts.cwd }

package/lib/chat/context.ts

@@ -49,7 +49,7 @@ export function formatTokens(n: number): string {
   return String(n);
 }
 
-type ModelUsage = Record<
+export type ModelUsage = Record<
   string,
   {
     inputTokens?: number;

package/lib/chat/runner.ts

@@ -5,6 +5,7 @@
  * events to whoever is watching. Restarting the server leaves turns running.
  */
 
+import { handoffAfterPlan } from "../models/plan-step";
 import { randomUUID } from "crypto";
 import os from "os";
 import { db, type Session } from "../db";
@@ -163,6 +164,10 @@ function onWorkerEvent(sessionId: string, live: Live, e: WorkerEvent): void {
     clearTimeout(live.retiring);
     live.retiring = undefined;
     retireIfStale(sessionId, live);
+    if (e.state === "idle")
+      void handoffAfterPlan(sessionId, (id, m) => sendQueuedNow(id, m)).catch(
+        (err) => console.error(`Plan handoff for ${sessionId} failed:`, err)
+      );
   } else if (e.type === "suggestion") {
     emit(sessionId, e);
     if (live.retiring && e.text) retire(sessionId, live);

package/lib/chat/worker/extras.ts

@@ -1,6 +1,7 @@
 import type { Session } from "../../db";
 import type { ChatStartOptions } from "../driver";
 import { refreshSettings } from "../refresh";
+import { sessionProviderEnv } from "../../models/launch";
 
 export type RoleExtras = Pick<
   ChatStartOptions,
@@ -17,11 +18,22 @@ export type RoleExtras = Pick<
   // Start over on a new conversation once its context passes this many
   // tokens (lib/chat/refresh); absent or 0, never.
   refresh?: { atTokens: number; role: "orchestrator" | "task" };
+  // What the session's model provider needs in the agent's environment
+  // (lib/models), its key fetched just now. Held by this process only.
+  env?: Record<string, string>;
 };
 
 // What a session's role adds to its chat. An orchestrator's module loads
 // only for one; a task's chat gets its brief, as its terminal agent would.
 export async function roleExtras(session: Session): Promise<RoleExtras> {
+  const [extras, env] = await Promise.all([
+    ownExtras(session),
+    sessionProviderEnv(session),
+  ]);
+  return env ? { ...extras, env } : extras;
+}
+
+async function ownExtras(session: Session): Promise<RoleExtras> {
   if (session.role === "orchestrator")
     return (await import("../../orchestrator/worker")).orchestratorExtras(
       session

package/lib/chat/worker/host.ts

@@ -76,6 +76,8 @@ export class ChatHost {
   // resumed conversation's saved totals.
   private usage: UsageTotals | null;
   private turnsRecorded = 0;
+  // When the running turn started, for its wall time.
+  private turnStartedAt: number | null = null;
   // A permission mode change still on its way to the agent: a message sent
   // right after it (carrying out a plan) must not overtake it.
   private modeChange: Promise<void> = Promise.resolve();
@@ -161,7 +163,7 @@ export class ChatHost {
       throw new Error(`${session.agent_type} can't enforce this role's tools`);
     this.atRest = driver.atRest !== false;
     const cwd = session.working_directory.replace(/^~/, os.homedir());
-    const env = agentEnv(session.id);
+    const env = { ...agentEnv(session.id), ...extras.env };
     // Every chat can show visuals, asked for like any tool under its access
     // setting (they read files and reach the web). Not an orchestrator: it
     // reads other sessions' text, and a browser would be a way out for it.
@@ -290,7 +292,15 @@ export class ChatHost {
           this.setSuggestion(e.text);
           this.retireIfDone();
         } else if (e.type === "usage") {
-          this.usage = recordTurn(this.session, this.usage, e.totals);
+          const at = Date.now();
+          this.usage = recordTurn(
+            this.session,
+            this.usage,
+            e.totals,
+            at,
+            this.turnStartedAt === null ? null : at - this.turnStartedAt
+          );
+          this.turnStartedAt = null;
           this.turnsRecorded++;
         } else if (e.type === "usage_start") {
           // Too late once a turn was measured without it.
@@ -306,6 +316,7 @@ export class ChatHost {
           this.refreshIfFull(e.context.usedTokens);
         } else if (e.type === "turn_start") {
           this.turnsStarted++;
+          this.turnStartedAt ??= Date.now();
           if (this.behind) this.behind--;
           clearTimeout(this.restTimer);
           this.restTimer = undefined;

package/lib/chat/worker/main.ts

@@ -29,13 +29,16 @@ async function main(sessionId: string) {
   const { db } = await import("../../db");
   const { ChatHost } = await import("./host");
 
-  const session = db
+  const row = db
     .prepare(`SELECT * FROM sessions WHERE id = ?`)
     .get(sessionId) as Session | undefined;
-  if (!session) {
+  if (!row) {
     console.error(`No session ${sessionId}`);
     process.exit(1);
   }
+  const session = await (
+    await import("../../models/launch")
+  ).routeOrchestrator(row);
 
   const clients = new Set<net.Socket>();
   const emit = (e: WorkerEvent) => {

package/lib/db/migrations.ts

@@ -1400,6 +1400,34 @@ const migrations: Migration[] = [
       `);
     },
   },
+  {
+    id: 61,
+    name: "model_routing",
+    // One transaction, so a crash part-way leaves nothing to skip as
+    // "duplicate column" on the next start.
+    up: (db) =>
+      db.transaction(() => {
+        // What a session runs on beyond the model name: its provider (NULL is
+        // Anthropic), a task's own models by role (JSON), and a task's plan
+        // step (planning, then implementing). Each turn keeps the model,
+        // provider and role it ran as, how long it took, and the plan's 5-hour
+        // window when it ended, so spend compares by model. A check run keeps the
+        // provider its model came from.
+        for (const col of [
+          `ALTER TABLE sessions ADD COLUMN model_provider TEXT`,
+          `ALTER TABLE sessions ADD COLUMN model_overrides TEXT`,
+          `ALTER TABLE sessions ADD COLUMN plan_phase TEXT`,
+          `ALTER TABLE task_queue ADD COLUMN model_overrides TEXT`,
+          `ALTER TABLE chat_turns ADD COLUMN model TEXT`,
+          `ALTER TABLE chat_turns ADD COLUMN provider TEXT`,
+          `ALTER TABLE chat_turns ADD COLUMN role TEXT`,
+          `ALTER TABLE chat_turns ADD COLUMN duration_ms INTEGER`,
+          `ALTER TABLE chat_turns ADD COLUMN window_pct REAL`,
+          `ALTER TABLE check_runs ADD COLUMN provider TEXT`,
+        ])
+          db.exec(col);
+      })(),
+  },
 ];
 
 // `upTo`: stop after this id (tests that start from an older database).

package/lib/db/types.ts

@@ -14,6 +14,12 @@ export interface Session {
   parent_session_id: string | null;
   claude_session_id: string | null;
   model: string;
+  // The model's provider (lib/models), NULL for Anthropic.
+  model_provider?: string | null;
+  // A task's own models: "provider/model", or JSON of role to one.
+  model_overrides?: string | null;
+  // A task's plan step: planning on the planning model, then implementing.
+  plan_phase?: "planning" | "implementing" | null;
   system_prompt: string | null;
   group_path: string; // Deprecated - use project_id
   project_id: string | null;

package/lib/drafts.ts

@@ -23,6 +23,9 @@ export interface Draft {
   openPr: boolean;
   // A task that runs in a terminal rather than a chat (the default).
   terminal?: boolean;
+  // A task's own model, "provider/model"; unset, its roles' models
+  // (lib/models).
+  taskModel?: string | null;
   createdAt: number;
 }
 

package/lib/models/config.ts

@@ -0,0 +1,218 @@
+/**
+ * Model routing settings, kept in the settings table: the model each role
+ * runs on, the providers (the presets, edited or added to), the price table,
+ * and whether tasks plan first. A workspace may set its own role models.
+ * Nothing here is ever a key: a provider names where its key lives.
+ */
+
+import { z } from "zod";
+import { db } from "../db";
+import { ENV_NAME, RESERVED_ENV } from "../project-config/schema";
+import { KEY_ENV_NAME } from "./secrets";
+import {
+  DEFAULT_PRICES,
+  PROVIDER_PRESETS,
+  type ModelPrice,
+  type ModelProvider,
+} from "./providers";
+import {
+  ANTHROPIC,
+  DEFAULT_ROLE_MODELS,
+  MODEL_ROLES,
+  REF_PATTERN,
+  type ModelRole,
+} from "./roles";
+
+// A name passed to `tv` or `security` as its own argument: no option, no
+// line break, nothing a shell would read.
+const refName = z
+  .string()
+  .trim()
+  .min(1)
+  .max(120)
+  .regex(/^[\w .@:+-]+$/, "letters, digits, spaces and . @ : + - _ only")
+  .refine((s) => !s.startsWith("-"), "can't start with -");
+
+const secretRef = z.discriminatedUnion("kind", [
+  z.strictObject({
+    kind: z.literal("tokenvault"),
+    project: refName,
+    desc: refName.optional(),
+  }),
+  z.strictObject({
+    kind: z.literal("keychain"),
+    service: refName,
+    account: refName.optional(),
+  }),
+  z.strictObject({
+    kind: z.literal("env"),
+    name: z
+      .string()
+      .regex(
+        KEY_ENV_NAME,
+        "must be a *_API_KEY variable, not ANTHROPIC_, CLAUDE_ or AGENTOS_"
+      ),
+  }),
+]);
+
+const modelId = z
+  .string()
+  .trim()
+  .min(1)
+  .max(120)
+  .regex(/^[\w.:@/-]+$/, "not a model id");
+
+const baseUrl = z
+  .string()
+  .trim()
+  .url()
+  .refine(
+    (u) =>
+      u.startsWith("https://") ||
+      /^http:\/\/(localhost|127\.0\.0\.1|\[::1\])(:\d+)?(\/|$)/.test(u),
+    "must be https (or http on localhost)"
+  );
+
+const providerSchema = z.strictObject({
+  id: z
+    .string()
+    .regex(/^[a-z0-9][a-z0-9_-]{0,39}$/, "lowercase letters, digits, - and _"),
+  label: z.string().trim().min(1).max(60),
+  api: z.enum(["anthropic", "openai"]),
+  baseUrl: baseUrl.optional(),
+  auth: secretRef.optional(),
+  authHeader: z.enum(["bearer", "x-api-key"]).optional(),
+  models: z.array(modelId).max(40),
+  tiers: z
+    .strictObject({
+      opus: modelId.optional(),
+      sonnet: modelId.optional(),
+      haiku: modelId.optional(),
+    })
+    .optional(),
+  // A key can't hide here: no name an agent CLI reads for credentials.
+  env: z
+    .record(
+      z
+        .string()
+        .regex(ENV_NAME)
+        .refine((n) => !RESERVED_ENV.test(n), "is reserved"),
+      z.string().max(200)
+    )
+    .optional(),
+  enabled: z.boolean(),
+});
+
+const price = z.strictObject({
+  input: z.number().min(0),
+  cacheWrite: z.number().min(0),
+  cacheRead: z.number().min(0),
+  output: z.number().min(0),
+});
+
+const roleRef = z.string().trim().regex(REF_PATTERN, "must be provider/model");
+const roleMap = z.partialRecord(z.enum(MODEL_ROLES), roleRef);
+
+export const modelSettingsSchema = z.strictObject({
+  roles: roleMap.optional(),
+  providers: z.array(providerSchema).optional(),
+  prices: z.record(z.string().regex(REF_PATTERN), price).optional(),
+  planStep: z.boolean().optional(),
+});
+
+export const workspaceModelsSchema = z.strictObject({
+  roles: roleMap.optional(),
+});
+
+export type ModelSettings = z.infer<typeof modelSettingsSchema>;
+export type WorkspaceModels = z.infer<typeof workspaceModelsSchema>;
+export type RoleMap = Partial<Record<ModelRole, string>>;
+export const roleMapSchema = roleMap;
+
+const GLOBAL_KEY = "models";
+const workspaceKey = (id: string) => `models.workspace.${id}`;
+
+function read<T>(key: string, schema: z.ZodType<T>): T | null {
+  const row = db
+    .prepare(`SELECT value FROM settings WHERE key = ?`)
+    .get(key) as { value: string } | undefined;
+  if (!row) return null;
+  try {
+    const got = schema.safeParse(JSON.parse(row.value));
+    return got.success ? got.data : null;
+  } catch {
+    return null;
+  }
+}
+
+function problems(error: z.ZodError): string {
+  return error.issues
+    .map((i) => `${i.path.join(".") || "(top level)"}: ${i.message}`)
+    .join("; ");
+}
+
+function write<T>(key: string, schema: z.ZodType<T>, value: unknown): T {
+  const got = schema.safeParse(value ?? {});
+  if (!got.success)
+    throw new Error(`Invalid model settings: ${problems(got.error)}`);
+  db.prepare(
+    `INSERT INTO settings (key, value) VALUES (?, ?)
+     ON CONFLICT(key) DO UPDATE SET value = excluded.value`
+  ).run(key, JSON.stringify(got.data));
+  return got.data;
+}
+
+export const savedModelSettings = (): ModelSettings =>
+  read(GLOBAL_KEY, modelSettingsSchema) ?? {};
+
+export function setModelSettings(value: unknown): ModelSettings {
+  const parsed = modelSettingsSchema.safeParse(value ?? {});
+  if (!parsed.success)
+    throw new Error(`Invalid model settings: ${problems(parsed.error)}`);
+  const ids = parsed.data.providers?.map((p) => p.id) ?? [];
+  if (new Set(ids).size !== ids.length)
+    throw new Error("Invalid model settings: two providers share an id");
+  const anthropic = parsed.data.providers?.find((p) => p.id === ANTHROPIC);
+  if (anthropic && (anthropic.baseUrl || anthropic.auth || !anthropic.enabled))
+    throw new Error(
+      "Invalid model settings: anthropic is the Claude login; add another provider for an endpoint"
+    );
+  return write(GLOBAL_KEY, modelSettingsSchema, parsed.data);
+}
+
+export const workspaceModels = (workspaceId: string): WorkspaceModels =>
+  read(workspaceKey(workspaceId), workspaceModelsSchema) ?? {};
+
+export function setWorkspaceModels(
+  workspaceId: string,
+  value: unknown
+): WorkspaceModels {
+  const saved = write(workspaceKey(workspaceId), workspaceModelsSchema, value);
+  if (!Object.keys(saved.roles ?? {}).length)
+    db.prepare(`DELETE FROM settings WHERE key = ?`).run(
+      workspaceKey(workspaceId)
+    );
+  return saved;
+}
+
+// The settings in force: the presets with any saved edits laid over them
+// (by id), saved providers that aren't presets, and every role filled.
+export interface EffectiveModels {
+  roles: Record<ModelRole, string>;
+  providers: ModelProvider[];
+  prices: Record<string, ModelPrice>;
+  planStep: boolean;
+}
+
+export function effectiveModels(
+  saved: ModelSettings = savedModelSettings()
+): EffectiveModels {
+  const byId = new Map(PROVIDER_PRESETS.map((p) => [p.id, p]));
+  for (const p of saved.providers ?? []) byId.set(p.id, p as ModelProvider);
+  return {
+    roles: { ...DEFAULT_ROLE_MODELS, ...saved.roles },
+    providers: [...byId.values()],
+    prices: { ...DEFAULT_PRICES, ...saved.prices },
+    planStep: saved.planStep ?? true,
+  };
+}

package/lib/models/launch.ts

@@ -0,0 +1,93 @@
+/**
+ * What an agent is launched with to run on its routed model: the model
+ * name for its CLI, and the provider's endpoint and key in its environment.
+ */
+
+import type { Session } from "../db";
+import {
+  defaultRoute,
+  providerEnv,
+  providerFor,
+  providerShellEnv,
+  resolveRole,
+  setSessionModel,
+  type ModelOverride,
+} from "./resolve";
+import { ANTHROPIC, type ModelRole } from "./roles";
+import { redact, type SecretCommand } from "./secrets";
+
+const why = (error: unknown) =>
+  redact(error instanceof Error ? error.message : String(error));
+
+// A chat session's provider environment, or null on Anthropic.
+export async function sessionProviderEnv(
+  session: Pick<Session, "model" | "model_provider">,
+  exec?: SecretCommand
+): Promise<Record<string, string> | null> {
+  if (!session.model_provider) return null;
+  return providerEnv(providerFor(session.model_provider), session.model, exec);
+}
+
+// A terminal session's: plain settings for its tmux environment, and shell
+// text that sets the key from its reference inside the pane.
+export function sessionShellEnv(
+  session: Pick<Session, "model" | "model_provider">
+): { env: Record<string, string>; exportKey: string } {
+  if (!session.model_provider) return { env: {}, exportKey: "" };
+  return providerShellEnv(providerFor(session.model_provider), session.model);
+}
+
+export interface RoutedRun {
+  model: string;
+  provider: string;
+  env: Record<string, string>;
+}
+
+// A one-shot run (a review, a scope check, a title) on its role's model,
+// or on the role's default when its provider is off or its key can't be
+// read: a title or a review still happens.
+export async function routedRun(
+  role: ModelRole,
+  opts: { workspaceId?: string | null; override?: ModelOverride | null } = {},
+  exec?: SecretCommand
+): Promise<RoutedRun> {
+  try {
+    const routed = resolveRole(role, opts);
+    return {
+      model: routed.model,
+      provider: routed.provider.id,
+      env: await providerEnv(routed.provider, routed.model, exec),
+    };
+  } catch (error) {
+    const routed = defaultRoute(role);
+    console.warn(`[models] ${role}: ${why(error)}; on ${routed.ref} instead`);
+    return { model: routed.model, provider: ANTHROPIC, env: {} };
+  }
+}
+
+// An orchestrator runs on its role's model, as its workspace and Settings
+// say now: each worker start picks up a change. When that provider is off
+// or its key can't be read, it starts on the default rather than not at all.
+export async function routeOrchestrator(
+  session: Session,
+  exec?: SecretCommand
+): Promise<Session> {
+  if (session.role !== "orchestrator") return session;
+  let routed;
+  try {
+    routed = resolveRole("orchestrator", { workspaceId: session.workspace_id });
+    await providerEnv(routed.provider, routed.model, exec);
+  } catch (error) {
+    routed = defaultRoute("orchestrator");
+    console.warn(
+      `[models] orchestrator: ${why(error)}; on ${routed.ref} instead`
+    );
+  }
+  setSessionModel(session.id, routed);
+  return {
+    ...session,
+    model: routed.model,
+    model_provider:
+      routed.provider.id === ANTHROPIC ? null : routed.provider.id,
+  };
+}

package/lib/models/plan-step.ts

@@ -0,0 +1,94 @@
+/**
+ * A task's plan step: its first turn runs on the planning model and writes
+ * the implementation plan to a file in the worktree; when that turn ends,
+ * the same conversation switches to the implementation model (Claude Code
+ * keeps the transcript across a model change, and a resume takes a new
+ * model) and is told to follow the plan.
+ */
+
+import fs from "fs";
+import path from "path";
+import { db, type Session } from "../db";
+import { enqueue } from "../chat/queued";
+import { applySessionModel } from "./session-model";
+import {
+  resolveRole,
+  sessionWorkspace,
+  setSessionModel,
+  taskOverride,
+} from "./resolve";
+import { ANTHROPIC } from "./roles";
+
+export const PLAN_FILE = ".agentos/plan.md";
+
+export function planPrompt(task: string): string {
+  return `${task}
+
+---
+Plan first. This turn, only read and plan: write the implementation plan for the task above to \`${PLAN_FILE}\` in this worktree (what to change, file by file; the tests; how you'll verify it; anything ambiguous and the call you'd make). Change no other file, and don't commit, push or open a PR. End the turn once the plan is written: the implementation model takes over this conversation next and follows it.`;
+}
+
+export function implementPrompt(planned: boolean): string {
+  return planned
+    ? `The plan is in \`${PLAN_FILE}\`. Implement it now, finishing the task as its instructions above say. Where the plan turns out wrong, do what's right and say so.`
+    : `No plan was written to \`${PLAN_FILE}\`. Plan briefly, then implement the task, finishing it as its instructions above say.`;
+}
+
+// The worktree's plan folder, kept out of git by its own .gitignore.
+export function preparePlanDir(worktree: string): void {
+  const dir = path.join(worktree, path.dirname(PLAN_FILE));
+  fs.mkdirSync(dir, { recursive: true });
+  fs.writeFileSync(path.join(dir, ".gitignore"), "*\n");
+}
+
+export const implementMessageId = (sessionId: string) =>
+  `implement-${sessionId}`;
+
+// Called when a chat goes idle. A task whose plan turn has ended (it has a
+// recorded plan turn) moves to implementing, once, and the implementation
+// message is queued under a fixed id, which a restart sends rather than
+// repeats.
+export async function handoffAfterPlan(
+  sessionId: string,
+  send: (sessionId: string, id: string) => Promise<void>
+): Promise<boolean> {
+  const planned = db
+    .prepare(
+      `SELECT 1 FROM chat_turns WHERE session_id = ? AND role = 'plan' LIMIT 1`
+    )
+    .get(sessionId);
+  if (!planned) return false;
+  const s = db.prepare(`SELECT * FROM sessions WHERE id = ?`).get(sessionId) as
+    | Session
+    | undefined;
+  if (s?.plan_phase !== "planning" || !s.worktree_path) return false;
+  const routed = resolveRole("implement", {
+    workspaceId: sessionWorkspace(s),
+    override: taskOverride(s),
+  });
+  const id = implementMessageId(sessionId);
+  // The phase, the model and the message in one go: a restart part-way
+  // finds either none of them or all, and sends the message on the new model.
+  const claimed = db.transaction(() => {
+    const changed = db
+      .prepare(
+        `UPDATE sessions SET plan_phase = 'implementing'
+         WHERE id = ? AND plan_phase = 'planning' AND task_status = 'running'`
+      )
+      .run(sessionId).changes;
+    if (!changed) return false;
+    setSessionModel(sessionId, routed);
+    enqueue(sessionId, {
+      id,
+      text: implementPrompt(
+        fs.existsSync(path.join(s.worktree_path!, PLAN_FILE))
+      ),
+      origin: { kind: "decision", label: "Plan step" },
+    });
+    return true;
+  })();
+  if (!claimed) return false;
+  await applySessionModel(sessionId, s.model_provider ?? ANTHROPIC, routed);
+  await send(sessionId, id);
+  return true;
+}