arisa 4.3.4 → 5.0.2

package/src/core/agent/agent-manager.js

@@ -1,6 +1,6 @@
 import path from "node:path";
 import { readFile, stat, unlink } from "node:fs/promises";
-import { createAgentSession, DefaultResourceLoader, SessionManager, defineTool } from "@earendil-works/pi-coding-agent";
+import { createAgentSession, DefaultResourceLoader, SessionManager, SettingsManager, defineTool } from "@earendil-works/pi-coding-agent";
 import { Type } from "@sinclair/typebox";
 import { createPiRuntime, hasProviderAuth } from "./pi-runtime.js";
 import { resolveChatModelSelection } from "./model-selection.js";
@@ -9,6 +9,7 @@ import { withTimeout } from "./prompt-timeout.js";
 import { buildPiToolPolicy, getCoreCodingTools } from "./core-tools.js";
 import { createSystemShellTool } from "./system-shell-tool.js";
 import { clampModelThinkingLevel } from "./pi-runtime.js";
+import { clampModelSpeed, createModelSpeedController } from "./model-speed.js";
 import { arisaHomeDir, getChatPiSessionsDir } from "../../runtime/paths.js";
 
 const piValidationTimeoutMs = 60_000;
@@ -24,6 +25,113 @@ const arisaToolNames = [
   "send_artifact"
 ];
 
+function messageText(content) {
+  if (typeof content === "string") return content.trim();
+  if (!Array.isArray(content)) return "";
+  return content
+    .filter((item) => item?.type === "text" && typeof item.text === "string")
+    .map((item) => item.text.trim())
+    .filter(Boolean)
+    .join("\n");
+}
+
+export function formatPortableSessionHistory(messages = []) {
+  return messages
+    .map((message) => {
+      const text = messageText(message?.content);
+      if (!text) return "";
+      const role = message.role === "assistant"
+        ? "Assistant"
+        : message.role === "user"
+          ? "User"
+          : (message.customType ? `Session memory (${message.customType})` : "Session context");
+      return `${role}:\n${text}`;
+    })
+    .filter(Boolean)
+    .join("\n\n");
+}
+
+const estimatedImageTokens = 1_200;
+
+function estimateContentTokens(content) {
+  if (typeof content === "string") return Math.ceil(content.length / 4);
+  if (!Array.isArray(content)) return 0;
+  const chars = content.reduce((total, item) => {
+    if (item?.type === "image") return total + estimatedImageTokens * 4;
+    if (typeof item?.text === "string") return total + item.text.length;
+    if (typeof item?.thinking === "string") return total + item.thinking.length;
+    if (item?.type === "toolCall") {
+      return total + String(item.name || "").length + JSON.stringify(item.arguments || {}).length;
+    }
+    return total;
+  }, 0);
+  return Math.ceil(chars / 4);
+}
+
+function estimateMessageTokens(message) {
+  if (["user", "assistant", "custom", "toolResult"].includes(message?.role)) {
+    return estimateContentTokens(message.content);
+  }
+  if (message?.role === "bashExecution") {
+    return Math.ceil((String(message.command || "").length + String(message.output || "").length) / 4);
+  }
+  if (["branchSummary", "compactionSummary"].includes(message?.role)) {
+    return Math.ceil(String(message.summary || "").length / 4);
+  }
+  return 0;
+}
+
+function summarizeRetainedContext(messages = []) {
+  const sizes = messages.map((message) => ({
+    role: message?.role,
+    tokens: estimateMessageTokens(message)
+  }));
+  const estimatedTokens = sizes.reduce((total, item) => total + item.tokens, 0);
+  const toolResultTokens = sizes
+    .filter((item) => item.role === "toolResult")
+    .reduce((total, item) => total + item.tokens, 0);
+  const largestMessageTokens = sizes.reduce((largest, item) => Math.max(largest, item.tokens), 0);
+  return {
+    messages: messages.length,
+    estimatedTokens,
+    toolResultPercent: estimatedTokens ? toolResultTokens / estimatedTokens * 100 : 0,
+    largestMessagePercent: estimatedTokens ? largestMessageTokens / estimatedTokens * 100 : 0
+  };
+}
+
+function closeAgentSession(session) {
+  if (session?.close) return session.close();
+  if (session?.dispose) return session.dispose();
+  return undefined;
+}
+
+export const defaultScheduledTaskListLimit = 50;
+export const maxScheduledTaskListLimit = 100;
+
+export function selectScheduledTasks(tasks = [], { status, limit = defaultScheduledTaskListLimit } = {}) {
+  const parsedLimit = Number(limit);
+  const resolvedLimit = Math.min(
+    Math.max(Number.isFinite(parsedLimit) ? Math.trunc(parsedLimit) : defaultScheduledTaskListLimit, 1),
+    maxScheduledTaskListLimit
+  );
+  const allTasks = Array.isArray(tasks) ? tasks : [];
+  const orderedTasks = status
+    ? [...allTasks].reverse()
+    : [
+        ...allTasks.filter((task) => task.status === "pending" || task.status === "running").reverse(),
+        ...allTasks.filter((task) => task.status !== "pending" && task.status !== "running").reverse()
+      ];
+  const visibleTasks = orderedTasks.slice(0, resolvedLimit);
+
+  return {
+    tasks: visibleTasks,
+    total: allTasks.length,
+    returned: visibleTasks.length,
+    limit: resolvedLimit,
+    truncated: visibleTasks.length < allTasks.length
+  };
+}
+
 function isLocalBaseUrl(value) {
   if (typeof value !== "string" || !value.trim()) return false;
   try {
@@ -90,11 +198,16 @@ async function assertDirectory(dir, label) {
   }
 }
 
-async function createArisaResourceLoader({ cwd, agentDir }) {
+export function createPiSettingsManager(config) {
+  return SettingsManager.inMemory({ compaction: { ...config.pi.compaction } });
+}
+
+async function createArisaResourceLoader({ cwd, agentDir, settingsManager }) {
   const arisaAgentsContent = await readFile(arisaAgentsFile, "utf8");
   const resourceLoader = new DefaultResourceLoader({
     cwd,
     agentDir,
+    settingsManager,
     agentsFilesOverride: (current) => appendArisaAgentsFile(current, arisaAgentsContent)
   });
   await resourceLoader.reload();
@@ -110,22 +223,104 @@ export class AgentManager {
     this.logger = logger;
     this.sessions = new Map();
     this.pendingNewSessions = new Set();
+    this.pendingSessionHandoffs = new Map();
+    this.artifactDeliveryHandler = null;
+    this.sessionClosePromises = new Map();
+  }
+
+  setArtifactDeliveryHandler(handler) {
+    this.artifactDeliveryHandler = handler;
+  }
+
+  async deliverArtifact(payload) {
+    if (!this.artifactDeliveryHandler) throw new Error("Telegram artifact delivery is unavailable");
+    return this.artifactDeliveryHandler(payload);
+  }
+
+  closeCachedSession(sessionKey) {
+    const key = String(sessionKey);
+    const existing = this.sessions.get(key);
+    this.sessions.delete(key);
+    const closeSession = (existing?.session?.close || existing?.session?.dispose)
+      ? () => closeAgentSession(existing.session)
+      : null;
+    if (!closeSession) {
+      return this.sessionClosePromises.get(key) || Promise.resolve();
+    }
+
+    const previousClose = this.sessionClosePromises.get(key);
+    const closePromise = Promise.resolve(previousClose)
+      .catch(() => {})
+      .then(closeSession)
+      .catch((error) => {
+        this.logger?.error?.("agent", `session close failed for chat ${key}: ${error instanceof Error ? error.message : String(error)}`);
+      })
+      .finally(() => {
+        if (this.sessionClosePromises.get(key) === closePromise) {
+          this.sessionClosePromises.delete(key);
+        }
+      });
+    this.sessionClosePromises.set(key, closePromise);
+    return closePromise;
+  }
+
+  async waitForSessionClose(sessionKey) {
+    const key = String(sessionKey);
+    let closing = this.sessionClosePromises.get(key);
+    while (closing) {
+      await closing;
+      closing = this.sessionClosePromises.get(key);
+    }
   }
 
   setConfig(config) {
+    for (const key of this.sessions.keys()) this.closeCachedSession(key);
     this.config = config;
-    this.sessions.clear();
     this.pendingNewSessions.clear();
+    this.pendingSessionHandoffs.clear();
   }
 
-  resetSession(chatId) {
+  resetSession(chatId, { handoff = "", parentSession = "" } = {}) {
     const sessionKey = String(chatId);
-    this.sessions.delete(sessionKey);
+    this.closeCachedSession(sessionKey);
     this.pendingNewSessions.add(sessionKey);
+    const text = String(handoff || "").trim();
+    const parent = String(parentSession || "").trim();
+    if (text || parent) {
+      this.pendingSessionHandoffs.set(sessionKey, { text, parentSession: parent });
+    } else {
+      this.pendingSessionHandoffs.delete(sessionKey);
+    }
   }
 
   clearSessionCache(chatId) {
-    this.sessions.delete(String(chatId));
+    this.closeCachedSession(String(chatId));
+  }
+
+  async getRuntimeDiagnostic({ contextInspectionTimeoutMs } = {}) {
+    const contexts = await Promise.all([...this.sessions.entries()].map(async ([chatId, context]) => {
+      const base = { chatId };
+      try {
+        const stats = context.session.getSessionStats();
+        const retained = summarizeRetainedContext(context.session.messages);
+        return {
+          ...base,
+          ...retained,
+          tokens: stats.contextUsage?.tokens ?? null,
+          contextWindow: stats.contextUsage?.contextWindow ?? null,
+          percent: stats.contextUsage?.percent ?? null
+        };
+      } catch (error) {
+        return { ...base, error: error instanceof Error ? error.message : String(error) };
+      }
+    }));
+    return {
+      harness: "pi",
+      sessions: this.sessions.size,
+      closingSessions: this.sessionClosePromises.size,
+      managedProcessIds: [],
+      contexts
+    };
   }
 
   createSessionManager(chatId, workspaceDir = arisaInstallDir, sessionRevision = 0) {
@@ -133,36 +328,60 @@ export class AgentManager {
     const sessionDir = getChatPiSessionsDir(sessionKey, sessionRevision);
     if (this.pendingNewSessions.has(sessionKey)) {
       this.logger?.log("agent", `starting new persisted session for chat ${sessionKey}`);
-      return { sessionManager: SessionManager.create(workspaceDir, sessionDir), isNewSession: true };
+      const handoff = this.pendingSessionHandoffs.get(sessionKey);
+      const sessionManager = SessionManager.create(
+        workspaceDir,
+        sessionDir,
+        handoff?.parentSession ? { parentSession: handoff.parentSession } : undefined
+      );
+      if (handoff?.text) {
+        sessionManager.appendCustomMessageEntry(
+          "arisa-session-handoff",
+          handoff.text,
+          false,
+          { source: "telegram-new" }
+        );
+      }
+      return { sessionManager, isNewSession: true };
     }
     this.logger?.log("agent", `recovering persisted session for chat ${sessionKey}`);
     return { sessionManager: SessionManager.continueRecent(workspaceDir, sessionDir), isNewSession: false };
   }
 
-  async validatePiAgent() {
+  async validatePiAgent(config = this.config) {
     this.logger?.log("agent", "validating Pi session");
     const { authStorage, modelRegistry } = createPiRuntime({
-      provider: this.config.pi.provider,
-      apiKey: this.config.pi.apiKey
+      provider: config.pi.provider,
+      apiKey: config.pi.apiKey
     });
-    const model = modelRegistry.find(this.config.pi.provider, this.config.pi.model);
+    const model = modelRegistry.find(config.pi.provider, config.pi.model);
     if (!model) {
-      throw new Error(`Model not found: ${this.config.pi.provider}/${this.config.pi.model}`);
+      throw new Error(`Model not found: ${config.pi.provider}/${config.pi.model}`);
     }
-    if (requiresProviderAuth(model) && !this.config.pi.apiKey && !hasProviderAuth(this.config.pi.provider, { authStorage, modelRegistry })) {
-      throw new Error(`No auth found for ${this.config.pi.provider}. Provide a Pi API key in bootstrap, or authenticate with Pi login for this provider during bootstrap.`);
+    if (requiresProviderAuth(model) && !config.pi.apiKey && !hasProviderAuth(config.pi.provider, { authStorage, modelRegistry })) {
+      throw new Error(`No auth found for ${config.pi.provider}. Provide a Pi API key in bootstrap, or authenticate with Pi login for this provider during bootstrap.`);
     }
 
+    const settingsManager = createPiSettingsManager(config);
     const { session } = await createAgentSession({
       authStorage,
       modelRegistry,
       model,
-      sessionManager: SessionManager.inMemory(),
-    });
-    await withTimeout(promptAndThrowOnAssistantError(session, "Reply with exactly: OK"), {
-      timeoutMs: piValidationTimeoutMs,
-      label: "Pi validation prompt"
+      settingsManager,
+      sessionManager: SessionManager.inMemory()
     });
+    try {
+      await withTimeout(promptAndThrowOnAssistantError(session, "Reply with exactly: OK"), {
+        timeoutMs: piValidationTimeoutMs,
+        label: "Pi validation prompt"
+      });
+    } finally {
+      session.dispose();
+    }
+  }
+
+  async validateAgent(config = this.config) {
+    return this.validatePiAgent(config);
   }
 
   async getSessionContext(chatId, telegram) {
@@ -178,11 +397,16 @@ export class AgentManager {
           this.logger?.log("agent", `updating effort for chat ${sessionKey}: ${existing.session.thinkingLevel} -> ${desiredThinkingLevel}`);
           existing.session.setThinkingLevel(desiredThinkingLevel);
         }
+        const desiredSpeed = clampModelSpeed(existing.session.model, modelSelection.speed);
+        if (existing.speedController.speed !== desiredSpeed) {
+          this.logger?.log("agent", `updating speed for chat ${sessionKey}: ${existing.speedController.speed}x -> ${desiredSpeed}x`);
+          existing.speedController.setSpeed(desiredSpeed);
+        }
         this.logger?.log("agent", `reusing session for chat ${sessionKey}`);
         return existing;
       }
       this.logger?.log("agent", `model changed for chat ${sessionKey}: ${existing?.modelKey || "unknown"} -> ${effectiveModelKey}; recreating session`);
-      this.sessions.delete(sessionKey);
+      this.closeCachedSession(sessionKey);
       this.pendingNewSessions.add(sessionKey);
     }
 
@@ -196,6 +420,7 @@ export class AgentManager {
       throw new Error(`No auth found for ${this.config.pi.provider}. Re-run bootstrap and complete login for this provider before Telegram starts.`);
     }
     const thinkingLevel = clampModelThinkingLevel(model, modelSelection.thinkingLevel);
+    const speed = clampModelSpeed(model, modelSelection.speed);
 
     const policy = buildPiToolPolicy({
       config: this.config,
@@ -208,14 +433,16 @@ export class AgentManager {
       modelSelection.sessionRevision
     );
     const hasExistingSession = sessionManager.buildSessionContext().messages.length > 0;
-    this.logger?.log("agent", `${hasExistingSession ? "resuming" : "creating"} session for chat ${sessionKey} with model ${effectiveModelId} effort ${thinkingLevel}`);
+    this.logger?.log("agent", `${hasExistingSession ? "resuming" : "creating"} session for chat ${sessionKey} with model ${effectiveModelId} effort ${thinkingLevel} speed ${speed}x`);
     const customTools = [
       ...this.createTools(telegram, chatId, policy),
       createSystemShellTool({ workspaceDir: policy.workspaceDir, shell: policy.shell })
     ];
+    const settingsManager = createPiSettingsManager(this.config);
     const resourceLoader = await createArisaResourceLoader({
       cwd: policy.workspaceDir,
-      agentDir: arisaHomeDir
+      agentDir: arisaHomeDir,
+      settingsManager
     });
     const { session } = await createAgentSession({
       cwd: policy.workspaceDir,
@@ -228,8 +455,11 @@ export class AgentManager {
       tools: policy.tools,
       excludeTools: policy.excludeTools,
       customTools,
+      settingsManager,
       sessionManager
     });
+    const speedController = createModelSpeedController(session.agent.streamFn, speed);
+    session.agent.streamFn = speedController.streamFn;
 
     if (!hasExistingSession) {
       this.logger?.log("agent", `created new session for chat ${sessionKey}`);
@@ -239,14 +469,38 @@ export class AgentManager {
       })}`);
     }
 
-    const ctx = { session, modelId: effectiveModelId, modelKey: effectiveModelKey };
+    const ctx = { session, modelId: effectiveModelId, modelKey: effectiveModelKey, speedController };
     this.sessions.set(sessionKey, ctx);
     if (isNewSession) {
       this.pendingNewSessions.delete(sessionKey);
+      this.pendingSessionHandoffs.delete(sessionKey);
     }
     return ctx;
   }
 
+  async getAvailableModels(chatId) {
+    const { listProviderModels } = await import("./pi-runtime.js");
+    const runtime = createPiRuntime({ provider: this.config.pi.provider, apiKey: this.config.pi.apiKey });
+    return listProviderModels(this.config.pi.provider, runtime);
+  }
+
+  async setModelSpeed(chatId, speed) {
+    const context = this.sessions.get(String(chatId));
+    if (!context) return speed;
+    const effectiveSpeed = clampModelSpeed(context.session.model, speed);
+    context.speedController.setSpeed(effectiveSpeed);
+    return effectiveSpeed;
+  }
+
+  async close() {
+    const contexts = [...this.sessions.values()];
+    this.sessions.clear();
+    await Promise.allSettled([
+      ...this.sessionClosePromises.values(),
+      ...contexts.map((context) => closeAgentSession(context.session))
+    ]);
+  }
+
   async runTool({ name, request, chatId }) {
     await this.toolRegistry.load();
     this.logger?.log("agent", `run_tool ${name}`);
@@ -297,7 +551,7 @@ export class AgentManager {
       defineTool({
         name: "list_tools",
         label: "List tools",
-        description: "List Arisa core, native shell, and modular CLI tools with their capabilities.",
+        description: "List Arisa core, native shell, and modular CLI tools with their capabilities and daemon diagnostics.",
         parameters: Type.Object({}),
         execute: async () => {
           await this.toolRegistry.load();
@@ -313,7 +567,7 @@ export class AgentManager {
             shell: policy.shell.shellPath || (process.platform === "win32" ? "powershell" : "sh"),
             enabled: !(policy.excludeTools || []).includes("system_shell")
           }];
-          const cliTools = this.toolRegistry.list().map((tool) => ({
+          const cliTools = (await this.toolRegistry.listWithRuntime(chatId)).map((tool) => ({
             ...tool,
             source: "arisa-modular",
             invocation: "run_tool"
@@ -410,15 +664,20 @@ export class AgentManager {
       defineTool({
         name: "list_scheduled_tasks",
         label: "List scheduled tasks",
-        description: "List scheduled async tasks for the current Telegram chat.",
+        description: "List scheduled async tasks for the current Telegram chat. Results default to 50 tasks, always include pending/running tasks, and accept an optional limit up to 100.",
         parameters: Type.Object({
-          status: Type.Optional(Type.String())
+          status: Type.Optional(Type.String()),
+          limit: Type.Optional(Type.Integer({ minimum: 1, maximum: maxScheduledTaskListLimit }))
         }),
         execute: async (_id, params) => {
           const tasks = await this.taskStore.list({ chatId, status: params.status });
+          const result = selectScheduledTasks(tasks, {
+            status: params.status,
+            limit: params.limit
+          });
           return {
-            content: [{ type: "text", text: JSON.stringify(tasks, null, 2) }],
-            details: { tasks }
+            content: [{ type: "text", text: JSON.stringify(result, null, 2) }],
+            details: result
           };
         }
       }),

package/src/core/agent/auth-flow.js

@@ -1,10 +1,11 @@
 import { createPiRuntime, hasProviderAuth, supportsProviderOAuth } from "./pi-runtime.js";
+import { resolveChatModelSelection } from "./model-selection.js";
 
 const authInvalidatedPatterns = [
   /authentication token has been invalidated/i,
   /token (?:has been )?invalidated/i,
   /try signing in again/i,
-  /auth(?:entication)? token (?:expired|revoked|invalid)/i
+  /auth(?:entication)? token (?:is |has been )?(?:expired|revoked|invalid)/i
 ];
 
 const missingAuthPatterns = [
@@ -32,23 +33,26 @@ export function getPiAuthIssue(error) {
   return null;
 }
 
-export function getPiAuthStatus(config) {
+export function getPiAuthStatus(config, chatId = null) {
   const runtime = createPiRuntime({
     provider: config.pi.provider,
     apiKey: config.pi.apiKey
   });
+  const modelSelection = chatId == null
+    ? { provider: config.pi.provider, model: config.pi.model }
+    : resolveChatModelSelection(config, chatId);
 
   return {
-    provider: config.pi.provider,
-    model: config.pi.model,
+    provider: modelSelection.provider,
+    model: modelSelection.model,
     hasApiKey: Boolean(config.pi.apiKey),
     hasStoredAuth: hasProviderAuth(config.pi.provider, runtime),
     supportsOAuth: supportsProviderOAuth(config.pi.provider, runtime)
   };
 }
 
-export function buildPiAuthTelegramMessage({ config, issue = null, verified = false }) {
-  const status = getPiAuthStatus(config);
+export function buildPiAuthTelegramMessage({ config, chatId = null, issue = null, verified = false }) {
+  const status = getPiAuthStatus(config, chatId);
   let title = `Pi authentication status for ${status.provider}/${status.model}.`;
   if (issue) {
     title = `Pi authentication needs attention for ${status.provider}/${status.model}.`;
@@ -91,8 +95,8 @@ export function buildPiAuthTelegramMessage({ config, issue = null, verified = fa
   return lines.join("\n");
 }
 
-export function buildPiAuthRecoveryBlockedMessage({ config, issue = null, renewalActive = false }) {
-  const status = getPiAuthStatus(config);
+export function buildPiAuthRecoveryBlockedMessage({ config, chatId = null, issue = null, renewalActive = false }) {
+  const status = getPiAuthStatus(config, chatId);
   const lines = [
     `Pi authentication is not ready for ${status.provider}/${status.model}.`,
     "I did not send your message to the agent."

package/src/core/agent/model-selection.js

@@ -1,7 +1,13 @@
+import { normalizeModelSpeed } from "./model-speed.js";
+
 function chatKey(chatId) {
   return String(chatId);
 }
 
+export function getAgentConfig(config) {
+  return config.pi;
+}
+
 function normalizeSessionRevision(sessionRevision) {
   if (sessionRevision == null) return 0;
   if (!Number.isSafeInteger(sessionRevision) || sessionRevision < 0) {
@@ -11,12 +17,14 @@ function normalizeSessionRevision(sessionRevision) {
 }
 
 export function resolveChatModelSelection(config, chatId) {
-  const selection = config.pi.chatModels?.[chatKey(chatId)];
-  if (!selection || selection.provider !== config.pi.provider) {
+  const agentConfig = getAgentConfig(config);
+  const selection = agentConfig.chatModels?.[chatKey(chatId)];
+  if (!selection || selection.provider !== agentConfig.provider) {
     return {
-      provider: config.pi.provider,
-      model: config.pi.model,
-      thinkingLevel: config.pi.thinkingLevel,
+      provider: agentConfig.provider,
+      model: agentConfig.model,
+      thinkingLevel: agentConfig.thinkingLevel,
+      ...(agentConfig.speed !== undefined ? { speed: normalizeModelSpeed(agentConfig.speed) } : {}),
       sessionRevision: 0
     };
   }
@@ -24,7 +32,10 @@ export function resolveChatModelSelection(config, chatId) {
   return {
     provider: selection.provider,
     model: selection.model,
-    thinkingLevel: selection.thinkingLevel ?? config.pi.thinkingLevel,
+    thinkingLevel: selection.thinkingLevel ?? agentConfig.thinkingLevel,
+    ...(agentConfig.speed !== undefined
+      ? { speed: normalizeModelSpeed(selection.speed ?? agentConfig.speed) }
+      : {}),
     sessionRevision
   };
 }
@@ -37,29 +48,58 @@ export function resolveChatThinkingLevel(config, chatId) {
   return resolveChatModelSelection(config, chatId).thinkingLevel;
 }
 
-export function selectChatModel(config, chatId, model, { thinkingLevel } = {}) {
-  if (model.provider !== config.pi.provider) {
-    throw new Error(`Cannot select model from provider ${model.provider}; active provider is ${config.pi.provider}`);
+export function resolveChatSpeed(config, chatId) {
+  const speed = resolveChatModelSelection(config, chatId).speed;
+  if (speed === undefined) throw new Error("Model speed is not configured for the active runtime");
+  return speed;
+}
+
+export function selectChatModel(config, chatId, model, { thinkingLevel, speed } = {}) {
+  const agentConfig = getAgentConfig(config);
+  if (model.provider !== agentConfig.provider) {
+    throw new Error(`Cannot select model from provider ${model.provider}; active provider is ${agentConfig.provider}`);
   }
-  config.pi.chatModels ||= {};
+  agentConfig.chatModels ||= {};
   const key = chatKey(chatId);
-  const sessionRevision = (config.pi.chatModels[key]?.sessionRevision || 0) + 1;
-  config.pi.chatModels[key] = {
+  const sessionRevision = (agentConfig.chatModels[key]?.sessionRevision || 0) + 1;
+  agentConfig.chatModels[key] = {
     provider: model.provider,
     model: model.id,
     thinkingLevel,
+    ...(agentConfig.speed !== undefined
+      ? { speed: normalizeModelSpeed(speed ?? resolveChatModelSelection(config, chatId).speed) }
+      : {}),
     sessionRevision
   };
 }
 
 export function selectChatThinkingLevel(config, chatId, thinkingLevel) {
-  config.pi.chatModels ||= {};
+  const agentConfig = getAgentConfig(config);
+  agentConfig.chatModels ||= {};
   const key = chatKey(chatId);
   const current = resolveChatModelSelection(config, chatId);
-  config.pi.chatModels[key] = {
+  agentConfig.chatModels[key] = {
     provider: current.provider,
     model: current.model,
     thinkingLevel,
+    ...(current.speed !== undefined ? { speed: current.speed } : {}),
+    sessionRevision: current.sessionRevision
+  };
+}
+
+export function selectChatSpeed(config, chatId, speed) {
+  const agentConfig = getAgentConfig(config);
+  if (agentConfig.speed === undefined) {
+    throw new Error("Model speed is not configured for the active runtime");
+  }
+  agentConfig.chatModels ||= {};
+  const key = chatKey(chatId);
+  const current = resolveChatModelSelection(config, chatId);
+  agentConfig.chatModels[key] = {
+    provider: current.provider,
+    model: current.model,
+    thinkingLevel: current.thinkingLevel,
+    speed: normalizeModelSpeed(speed),
     sessionRevision: current.sessionRevision
   };
 }