2064 lines · 70.0 KB · TypeScript
RawDownload
import path from "path"import os from "os"import fs from "fs/promises"import z from "zod"import { Identifier } from "../id/id"import { MessageV2 } from "./message-v2"import { Log } from "../util/log"import { SessionRevert } from "./revert"import { Session } from "."import { Agent } from "../agent/agent"import { Provider } from "../provider/provider"import { type Tool as AITool, tool, jsonSchema, type ToolCallOptions } from "ai"import { SessionCompaction } from "./compaction"import { Instance } from "../project/instance"import { Bus } from "../bus"import { ProviderTransform } from "../provider/transform"import { SystemPrompt } from "./system"import { InstructionPrompt } from "./instruction"import { Plugin } from "../plugin"import PROMPT_PLAN from "../session/prompt/plan.txt"import PROMPT_WRITE from "../agent/prompt/write.txt"import PROMPT_ML from "../agent/prompt/ml.txt"import PROMPT_RESEARCH from "../agent/prompt/research.txt"import PROMPT_BIOLOGY from "../agent/prompt/biology.txt"import PROMPT_PHYSICS from "../agent/prompt/physics.txt"import BUILD_SWITCH from "../session/prompt/build-switch.txt"import MAX_STEPS from "../session/prompt/max-steps.txt"import { defer } from "../util/defer"import { clone } from "remeda"import { ToolRegistry } from "../tool/registry"import { MCP } from "../mcp"import { LSP } from "../lsp"import { ReadTool } from "../tool/read"import { ListTool } from "../tool/ls"import { FileTime } from "../file/time"import { Flag } from "../flag/flag"import { RSITrajectory } from "./rsi/trajectory"import { SessionReview } from "./review"import { RLMArtifacts } from "./rlm/artifacts"import { ulid } from "ulid"import { spawn } from "child_process"import { Command } from "../command"import { $, fileURLToPath } from "bun"import { ConfigMarkdown } from "../config/markdown"import { Config } from "../config/config"import { computeBillingMode } from "./billing-gate"import { SessionSummary } from "./summary"import { NamedError } from "@synsci/util/error"import { fn } from "@/util/fn"import { SessionProcessor } from "./processor"import { TaskTool } from "@/tool/task"import { Tool } from "@/tool/tool"import { PermissionNext } from "@/permission/next"import { SessionStatus } from "./status"import { LLM } from "./llm"import { iife } from "@/util/iife"import { correctImageMime } from "@/util/image"import { Shell } from "@/shell/shell"import { Truncate } from "@/tool/truncation"import { Memory } from "@/settings/memory" // @ts-ignoreglobalThis.AI_SDK_LOG_WARNINGS = false export namespace SessionPrompt {  const log = Log.create({ service: "session.prompt" })  export const OUTPUT_TOKEN_MAX = Flag.OPENSCIENCE_EXPERIMENTAL_OUTPUT_TOKEN_MAX || 32_000  const ARTIFACT_AGENTS = ["research", "biology", "ml"]  // Science agents that dispatch GPU/compute work and should honor billing.compute.  const COMPUTE_AGENTS = new Set(["research", "biology", "physics", "ml"])   const state = Instance.state(    () => {      const data: Record<        string,        {          abort: AbortController          callbacks: {            resolve(input: MessageV2.WithParts): void            reject(): void          }[]        }      > = {}      return data    },    async (current) => {      for (const item of Object.values(current)) {        item.abort.abort()        for (const callback of item.callbacks) {          callback.reject()        }      }    },  )   export function assertNotBusy(sessionID: string) {    const match = state()[sessionID]    if (match) throw new Session.BusyError(sessionID)  }   export const PromptInput = z.object({    sessionID: Identifier.schema("session"),    messageID: Identifier.schema("message").optional(),    model: z      .object({        providerID: z.string(),        modelID: z.string(),      })      .optional(),    agent: z.string().optional(),    noReply: z.boolean().optional(),    tools: z      .record(z.string(), z.boolean())      .optional()      .describe(        "@deprecated tools and permissions have been merged, you can set permissions on the session itself now",      ),    system: z.string().optional(),    variant: z.string().optional(),    tier: z.enum(["fast", "pro", "ultra"]).optional(),    // Per-request fast API param (gpt-5.5 → service_tier:"priority"). See llm.ts.    fast: z.boolean().optional(),    parts: z.array(      z.discriminatedUnion("type", [        MessageV2.TextPart.omit({          messageID: true,          sessionID: true,        })          .partial({            id: true,          })          .meta({            ref: "TextPartInput",          }),        MessageV2.FilePart.omit({          messageID: true,          sessionID: true,        })          .partial({            id: true,          })          .meta({            ref: "FilePartInput",          }),        MessageV2.AgentPart.omit({          messageID: true,          sessionID: true,        })          .partial({            id: true,          })          .meta({            ref: "AgentPartInput",          }),        MessageV2.SubtaskPart.omit({          messageID: true,          sessionID: true,        })          .partial({            id: true,          })          .meta({            ref: "SubtaskPartInput",          }),      ]),    ),  })  export type PromptInput = z.infer<typeof PromptInput>   export const prompt = fn(PromptInput, async (input) => {    const session = await Session.get(input.sessionID)    await SessionRevert.cleanup(session)     const message = await createUserMessage(input).catch((e) => {      // e.g. no providers are available at all — surface the failure to the      // session (the web UI listens for session.error) instead of only throwing.      const message = e instanceof Error ? e.message : String(e)      Bus.publish(Session.Event.Error, {        sessionID: input.sessionID,        error: new NamedError.Unknown({ message }).toObject(),      })      throw e    })    await Session.touch(input.sessionID)     // this is backwards compatibility for allowing `tools` to be specified when    // prompting    const permissions: PermissionNext.Ruleset = []    for (const [tool, enabled] of Object.entries(input.tools ?? {})) {      permissions.push({        permission: tool,        action: enabled ? "allow" : "deny",        pattern: "*",      })    }    if (permissions.length > 0) {      session.permission = permissions      await Session.update(session.id, (draft) => {        draft.permission = permissions      })    }     if (input.noReply === true) {      return message    }     return loop(input.sessionID)  })   export async function resolvePromptParts(template: string): Promise<PromptInput["parts"]> {    const parts: PromptInput["parts"] = [      {        type: "text",        text: template,      },    ]    const files = ConfigMarkdown.files(template)    const seen = new Set<string>()    await Promise.all(      files.map(async (match) => {        const name = match[1]        if (seen.has(name)) return        seen.add(name)        const filepath = name.startsWith("~/")          ? path.join(os.homedir(), name.slice(2))          : path.resolve(Instance.worktree, name)         const stats = await fs.stat(filepath).catch(() => undefined)        if (!stats) {          const agent = await Agent.get(name)          if (agent) {            parts.push({              type: "agent",              name: agent.name,            })          }          return        }         if (stats.isDirectory()) {          parts.push({            type: "file",            url: `file://${filepath}`,            filename: name,            mime: "application/x-directory",          })          return        }         parts.push({          type: "file",          url: `file://${filepath}`,          filename: name,          mime: "text/plain",        })      }),    )    return parts  }   function start(sessionID: string) {    const s = state()    if (s[sessionID]) return    const controller = new AbortController()    s[sessionID] = {      abort: controller,      callbacks: [],    }    return controller.signal  }   export function cancel(sessionID: string) {    log.info("cancel", { sessionID })    const s = state()    const match = s[sessionID]    if (!match) return    match.abort.abort()    for (const item of match.callbacks) {      item.reject()    }    delete s[sessionID]    SessionStatus.set(sessionID, { type: "idle" })    return  }   export const loop = fn(Identifier.schema("session"), async (sessionID) => {    const abort = start(sessionID)    if (!abort) {      return new Promise<MessageV2.WithParts>((resolve, reject) => {        const callbacks = state()[sessionID].callbacks        callbacks.push({ resolve, reject })      })    }     using _ = defer(() => cancel(sessionID))     let step = 0    const session = await Session.get(sessionID)    while (true) {      SessionStatus.set(sessionID, { type: "busy" })      log.info("loop", { step, sessionID })      if (abort.aborted) break      let msgs = await MessageV2.filterCompacted(MessageV2.stream(sessionID))       let lastUser: MessageV2.User | undefined      let lastAssistant: MessageV2.Assistant | undefined      let lastFinished: MessageV2.Assistant | undefined      let tasks: (MessageV2.CompactionPart | MessageV2.SubtaskPart)[] = []      for (let i = msgs.length - 1; i >= 0; i--) {        const msg = msgs[i]        if (!lastUser && msg.info.role === "user") lastUser = msg.info as MessageV2.User        if (!lastAssistant && msg.info.role === "assistant") lastAssistant = msg.info as MessageV2.Assistant        if (!lastFinished && msg.info.role === "assistant" && msg.info.finish)          lastFinished = msg.info as MessageV2.Assistant        if (lastUser && lastFinished) break        const task = msg.parts.filter((part) => part.type === "compaction" || part.type === "subtask")        if (task && !lastFinished) {          tasks.push(...task)        }      }       if (!lastUser) throw new Error("No user message found in stream. This should never happen.")      const bareMode = lastUser.tools?.["*"] === false      if (        lastAssistant?.finish &&        (!["tool-calls", "unknown"].includes(lastAssistant.finish) || bareMode) &&        lastUser.id < lastAssistant.id      ) {        log.info("exiting loop", { sessionID, bareMode })        // RSI: capture trajectory from ultra agent sessions (async, non-blocking)        if (lastUser.agent && RSITrajectory.ARTIFACT_AGENTS.includes(lastUser.agent as any)) {          RSITrajectory.pipeline(sessionID).catch(() => {})        }        // WS11 reviewer gate: blind review of the final answer, appended as a        // footer note. Off by default (config.experimental.reviewGate); the gate        // self-skips non-reviewable/trivial turns. Fire-and-forget, never blocks.        if (lastUser.agent) {          SessionReview.gate({ sessionID, agent: lastUser.agent, model: lastUser.model }).catch(() => {})        }        break      }       step++      if (step === 1)        ensureTitle({          session,          modelID: lastUser.model.modelID,          providerID: lastUser.model.providerID,          history: msgs,        }).catch((error) => log.error("failed to generate session title", { error }))       const model = await Provider.getModel(lastUser.model.providerID, lastUser.model.modelID).catch((e) => {        if (Provider.ModelNotFoundError.isInstance(e)) return undefined        throw e      })      // The requested model has no available provider (e.g. the API key was      // removed) — surface a session error instead of crashing the loop.      if (!model) {        const error = new NamedError.Unknown({          message: `Model ${lastUser.model.providerID}/${lastUser.model.modelID} is not available. Add your own API key (\`openscience keys add\`) or connect a managed account (\`openscience login\`), then choose a model.`,        }).toObject()        Bus.publish(Session.Event.Error, { sessionID, error })        await Session.updateMessage({          id: await MessageV2.nextMessageID(sessionID),          role: "assistant",          parentID: lastUser.id,          sessionID,          mode: lastUser.agent,          agent: lastUser.agent,          path: {            cwd: Instance.directory,            root: Instance.worktree,          },          cost: 0,          tokens: {            input: 0,            output: 0,            reasoning: 0,            cache: { read: 0, write: 0 },          },          modelID: lastUser.model.modelID,          providerID: lastUser.model.providerID,          error,          time: {            created: Date.now(),            completed: Date.now(),          },        })        break      }      const task = tasks.pop()       // pending subtask      // TODO: centralize "invoke tool" logic      if (task?.type === "subtask") {        const taskTool = await TaskTool.init()        const taskModel = task.model ? await Provider.getModel(task.model.providerID, task.model.modelID) : model        const assistantMessage = (await Session.updateMessage({          id: await MessageV2.nextMessageID(sessionID),          role: "assistant",          parentID: lastUser.id,          sessionID,          mode: task.agent,          agent: task.agent,          path: {            cwd: Instance.directory,            root: Instance.worktree,          },          cost: 0,          tokens: {            input: 0,            output: 0,            reasoning: 0,            cache: { read: 0, write: 0 },          },          modelID: taskModel.id,          providerID: taskModel.providerID,          time: {            created: Date.now(),          },        })) as MessageV2.Assistant        let part = (await Session.updatePart({          id: Identifier.ascending("part"),          messageID: assistantMessage.id,          sessionID: assistantMessage.sessionID,          type: "tool",          callID: ulid(),          tool: TaskTool.id,          state: {            status: "running",            input: {              prompt: task.prompt,              description: task.description,              subagent_type: task.agent,              command: task.command,            },            time: {              start: Date.now(),            },          },        })) as MessageV2.ToolPart        const taskArgs = {          prompt: task.prompt,          description: task.description,          subagent_type: task.agent,          command: task.command,        }        await Plugin.trigger(          "tool.execute.before",          {            tool: "task",            sessionID,            callID: part.id,          },          { args: taskArgs },        )        let executionError: Error | undefined        const taskAgent = await Agent.get(task.agent)        const taskCtx: Tool.Context = {          agent: task.agent,          messageID: assistantMessage.id,          sessionID: sessionID,          abort,          callID: part.callID,          extra: { bypassAgentCheck: true },          messages: msgs,          async metadata(input) {            await Session.updatePart({              ...part,              type: "tool",              state: {                ...part.state,                ...input,              },            } satisfies MessageV2.ToolPart)          },          async ask(req) {            await PermissionNext.ask({              ...req,              sessionID: sessionID,              ruleset: PermissionNext.merge(taskAgent.permission, session.permission ?? []),            })          },        }        const result = await taskTool.execute(taskArgs, taskCtx).catch((error) => {          executionError = error          log.error("subtask execution failed", { error, agent: task.agent, description: task.description })          return undefined        })        await Plugin.trigger(          "tool.execute.after",          {            tool: "task",            sessionID,            callID: part.id,          },          result,        )        assistantMessage.finish = "tool-calls"        assistantMessage.time.completed = Date.now()

Lines 1–500 of 2064.