806 lines · 29.4 KB · TypeScript
RawDownload
import { EventEmitter } from "node:events";import { randomUUID } from "node:crypto";import { closeSync, constants, fsyncSync, openSync, readSync, statSync } from "node:fs";import {  createAgentSessionFromServices,  createAgentSessionRuntime,  createAgentSessionServices,  getAgentDir,  SessionManager,  shouldCompact,  type AgentSessionEvent,  type AgentSessionRuntime,  type ExtensionUIContext,} from "@earendil-works/pi-coding-agent";import { Effect } from "effect";import type { AgentImageInput } from "../../../shared/agent/agent-image-input";import type { AgentQueueAction } from "../../../shared/agent/agent-turn";import {  applyRuntimeEnvInjections,  buildAgentSessionOptionsSync,  runtimeOptionsFingerprint,  resolveAgentCwdEffect,  type RuntimeStartOptions,} from "./pi-runtime-helpers";import {  refreshPiModels,  resolvePiModelSelection,  toPiThinkingLevel,} from "./pi-runtime-models";import { getProviderHub } from "./provider-hub";import { attachGoalDriver } from "./goal-driver";import { createGoalPromptExtension } from "./goal-prompt";import { findRuntimeSessionForLookup, piStatusFromEvents } from "./pi-runtime-state";import { configuredPiSessionDir, findSessionFile } from "./sessions-store";import { getGlobalSingleton } from "./instances";import { connectorsRevisionSync } from "./connectors-service";import type {  LoggedPiEvent,  PiAgentSession,  PiAgentStatus,  PiContextUsage,  PiPromptOptions,} from "./pi-runtime-types"; type PiEvent = LoggedPiEvent["event"]; function comparableQueuedText(text: string): string {  const marker = "\n\nUser prompt:\n";  const index = text.lastIndexOf(marker);  return (index === -1 ? text : text.slice(index + marker.length)).trim();} export function takeQueuedFollowUp(  followUp: readonly string[],  message: string,): { selected: string; before: string[]; after: string[] } | null {  const exactIndex = followUp.indexOf(message);  const target = comparableQueuedText(message);  const index =    exactIndex >= 0      ? exactIndex      : followUp.findIndex((candidate) => comparableQueuedText(candidate) === target);  if (index < 0) return null;  return {    selected: followUp[index]!,    before: followUp.slice(0, index),    after: followUp.slice(index + 1),  };} export function planQueuedFollowUpMutation(  followUp: readonly string[],  message: string,  action: AgentQueueAction,  replacement?: string,): { promoted: string | null; followUp: string[] } | null {  const selected = takeQueuedFollowUp(followUp, message);  if (!selected) return null;  if (action === "replace" && !replacement) {    throw new Error("Replacement text is required.");  }  return {    promoted: action === "promote" ? selected.selected : null,    followUp:      action === "replace"        ? [...selected.before, replacement!, ...selected.after]        : [...selected.before, ...selected.after],  };} type QueueTransport = {  steer: (message: string, images?: AgentImageInput[]) => Promise<void>;  followUp: (message: string, images?: AgentImageInput[]) => Promise<void>;}; export async function restoreQueuedMessages(  session: QueueTransport,  cleared: { steering: readonly string[]; followUp: readonly string[] },  mutation: { promoted: string | null; followUp: readonly string[] } | null,  images: AgentImageInput[] = [],): Promise<void> {  for (const queued of cleared.steering) await session.steer(queued);  if (mutation?.promoted) await session.steer(mutation.promoted, images);  for (const queued of mutation?.followUp ?? cleared.followUp) await session.followUp(queued);}  /** Appended to the system prompt for vision-capable models. Kept as an extra *  section rather than a replacement so first-party extensions still apply. */const VISION_GUIDANCE =  "When an image is attached, inspect it carefully before answering. State only details visible in the image. Never invent labels, UI elements, text, or facts. Say when details are too small or uncertain. Give a concise answer. Use available tools to inspect supplied files when helpful.";  export function selectPiRuntimeModel(  models: Awaited<ReturnType<typeof refreshPiModels>>["models"],  requestedModelId: string,) {  const exact = models.find((model) => model.id === requestedModelId);  if (exact) return exact;  const separator = requestedModelId.indexOf("/");  if (separator > 0) {    const providerId = requestedModelId.slice(0, separator);    const rawId = requestedModelId.slice(separator + 1);    const qualified = models.filter(      (model) => model.providerId === providerId && (model.rawId === rawId || model.id === rawId),    );    if (qualified.length === 1) return qualified[0];    if (qualified.length > 1) throw new Error(`Model '${requestedModelId}' is ambiguous.`);  }  const unqualified = models.filter(    (model) => model.rawId === requestedModelId || model.name === requestedModelId,  );  if (unqualified.length === 1) return unqualified[0];  if (unqualified.length > 1) throw new Error(`Model '${requestedModelId}' is ambiguous.`);  return null;} export function resolvePiRuntimeStartOptions(  current: RuntimeStartOptions,  running: boolean,  requested?: RuntimeStartOptions,): RuntimeStartOptions {  return structuredClone(requested ?? (running ? current : {}));} function runtimeFingerprint(  modelId: string,  cwd: string,  piSessionId: string | null,  options: RuntimeStartOptions,) {  return JSON.stringify({    modelId,    cwd,    piSessionId: piSessionId ?? "",    options: runtimeOptionsFingerprint(options),    connectors: connectorsRevisionSync(),  });} export function shouldRestartAfterPromptError(error: unknown): boolean {  return (    error instanceof Error && /Cannot continue from message role: assistant/i.test(error.message)  );} type PiResourceDiagnostic = {  type: "info" | "warning" | "error";  message: string;  path?: string;}; function diagnosticsMap(): Map<string, PiResourceDiagnostic[]> {  return getGlobalSingleton(    "piResourceDiagnostics",    () => new Map<string, PiResourceDiagnostic[]>(),  );} export function piResourceDiagnostics(agentDir?: string): PiResourceDiagnostic[] {  const map = diagnosticsMap();  if (agentDir) return map.get(agentDir) ?? [];  return [...map.values()].flat();} class PiSdkSession extends EventEmitter implements PiAgentSession {  private runtime: AgentSessionRuntime | null = null;  private unsubscribe: (() => void) | null = null;  private eventSeq = 0;  private eventLog: LoggedPiEvent[] = [];  private activePromptCount = 0;  private lastError: string | null = null;  private currentFingerprint = "";  private currentPiSessionId: string | null = null;  private currentCwd = "";  private currentModelId = "";  private currentStartOptions: RuntimeStartOptions = {};  private agentDir = "";  private queueEventBufferDepth = 0;  private bufferedQueueEvent: PiEvent | null = null;  private extensionUiPending = new Map<    string,    { method: "select" | "confirm" | "input" | "editor"; resolve: (value: unknown) => void }  >();   ensureStarted(    modelId: string,    cwd?: string,    piSessionId?: string | null,    options?: RuntimeStartOptions,  ): Promise<void> {    const effectiveOptions = resolvePiRuntimeStartOptions(      this.currentStartOptions,      Boolean(this.runtime),      options,    );    return Effect.runPromise(this.ensureStartedEffect(modelId, cwd, piSessionId, effectiveOptions));  }   private ensureStartedEffect(    modelId: string,    cwd: string | undefined,    piSessionId: string | null | undefined,    options: RuntimeStartOptions,  ): Effect.Effect<void, unknown> {    return Effect.gen(      function* (this: PiSdkSession) {        const resolvedCwd = yield* resolveAgentCwdEffect(cwd);        const desiredSessionId = piSessionId ?? null;        const fingerprint = runtimeFingerprint(modelId, resolvedCwd, desiredSessionId, options);        if (this.runtime && this.currentFingerprint === fingerprint) return;         yield* this.stopEffect();        this.eventSeq = 0;        this.eventLog = [];        this.activePromptCount = 0;        this.lastError = null;         const { models } = yield* Effect.tryPromise({          try: () => refreshPiModels(),          catch: (error) => error,        });        const selectedModel = selectPiRuntimeModel(models, modelId);        if (!selectedModel) {          return yield* Effect.fail(            new Error(`Model '${modelId}' is not available from /v1/models.`),          );        }        const resolvedSelection = resolvePiModelSelection(selectedModel.id);        const providerId = selectedModel.providerId ?? resolvedSelection.providerId;        const backendModelId = selectedModel.rawId ?? resolvedSelection.modelId;         // One shared ModelRuntime across sessions and the provider hub: a        // sign-in completed in settings is live for the next turn.        const sharedModelRuntime = yield* Effect.tryPromise({          try: () => getProviderHub(),          catch: (error) => error,        });         const sessionOptions = buildAgentSessionOptionsSync({ options, cwd: resolvedCwd });        applyRuntimeEnvInjections(sessionOptions.envInjections);        // Expose the current session's model so the automations extension can        // default a scheduled run to the same model the user is talking to.        applyRuntimeEnvInjections({ LOCAL_STUDIO_MODEL_ID: modelId });        const sessionDir = configuredPiSessionDir(resolvedCwd);        const resumeFile = desiredSessionId ? findSessionFile(resolvedCwd, desiredSessionId) : null;        const sessionManager = resumeFile          ? SessionManager.open(resumeFile, sessionDir, resolvedCwd)          : SessionManager.create(resolvedCwd, sessionDir);        const resuming = Boolean(resumeFile);        const agentDir = getAgentDir();        const extensionUiContext = this.extensionUiContext();        const recordExtensionEvent = (event: PiEvent) => this.recordEvent(event);        const runtime = yield* Effect.tryPromise({          try: () =>            createAgentSessionRuntime(              ({ cwd, agentDir, sessionManager, sessionStartEvent }) =>                Effect.runPromise(                  Effect.gen(function* () {                    const services = yield* Effect.tryPromise({                      try: () =>                        createAgentSessionServices({                          cwd,                          agentDir,                          modelRuntime: sharedModelRuntime,                          resourceLoaderOptions: {                            additionalSkillPaths: sessionOptions.skills,                            additionalExtensionPaths: sessionOptions.extensionPaths,                            additionalPromptTemplatePaths: sessionOptions.promptTemplatePaths,                            // In-process: the goal section is injected per turn                            // via before_agent_start, keyed by the canonical                            // piSessionId this SessionManager owns. Runs here so                            // it never depends on the RPC extension's session id                            // (which differs and left the goal unread — #284).                            extensionFactories: [                              {                                name: "local-studio-goal",                                factory: createGoalPromptExtension(() =>                                  sessionManager.getSessionId(),                                ),                              },                            ],                            // Vision guidance is APPENDED, not substituted. This branch used to                            // set noExtensions/noSkills/noContextFiles and replace the whole                            // system prompt, which silently disabled every first-party extension                            // (session goal, artifact policy, subagents) on any                            // vision-capable model — i.e. on the primary model.                            ...(selectedModel.vision                              ? {                                  appendSystemPromptOverride: (base: string[]) => [                                    ...base,                                    VISION_GUIDANCE,                                  ],                                }                              : {}),                          },                        }),                      catch: (error) => error,                    });                    const model = services.modelRuntime.getModel(providerId, backendModelId);                    if (!model) {                      return yield* Effect.fail(                        new Error(                          `Model '${providerId}/${backendModelId}' is not available to the SDK runtime.`,                        ),                      );                    }                    const created = yield* Effect.tryPromise({                      try: () =>                        createAgentSessionFromServices({                          services,                          sessionManager,                          sessionStartEvent,                          model,                          thinkingLevel: selectedModel.reasoning                            ? toPiThinkingLevel(options.thinkingLevel ?? "high")                            : undefined,                        }),                      catch: (error) => error,                    });                    const activeToolNames =                      options.toolAccess === "read_only"                        ? ["read", "grep", "find", "ls"]                        : created.session.getAllTools().map((tool) => tool.name);                    created.session.setActiveToolsByName(activeToolNames);                    yield* Effect.tryPromise({                      try: () =>                        created.session.bindExtensions({                          mode: "rpc",                          uiContext: extensionUiContext,                          onError: (error) => {                            recordExtensionEvent({                              type: "extension_error",                              error: error.error,                              extensionPath: error.extensionPath,                              event: error.event,                            });                          },                        }),                      catch: (error) => error,                    });                    const extensionErrors = services.resourceLoader                      .getExtensions()                      .errors.map(({ path, error }) => ({                        type: "error" as const,                        message: `Failed to load extension "${path}": ${error}`,                        path,                      }));                    const diagnostics = [...services.diagnostics, ...extensionErrors];                    diagnosticsMap().set(                      agentDir,                      diagnostics.map((d) => ({                        type: d.type as PiResourceDiagnostic["type"],                        message: d.message,                        path: "path" in d ? (d as { path?: string }).path : undefined,                      })),                    );                    return {                      ...created,                      services,                      diagnostics,                    };                  }),                ),              {                cwd: resolvedCwd,                agentDir,                sessionManager,                sessionStartEvent: {                  type: "session_start",                  reason: resuming ? "resume" : "startup",                },              },            ),          catch: (error) => error,        });         this.runtime = runtime;        this.agentDir = agentDir;        this.currentModelId = modelId;        this.currentCwd = resolvedCwd;        this.currentPiSessionId = runtime.session.sessionId || desiredSessionId;        this.currentFingerprint = fingerprint;        this.currentStartOptions = options;        this.unsubscribe = runtime.session.subscribe((event) => this.recordEvent(event));      }.bind(this),    );  }   prompt(    message: string,    onEvent: (event: PiEvent, seq: number) => void,    options: PiPromptOptions = {},  ): Promise<void> {    return Effect.runPromise(this.promptEffect(message, onEvent, options));  }   private promptEffect(    message: string,    onEvent: (event: PiEvent, seq: number) => void,    options: PiPromptOptions,  ): Effect.Effect<void, unknown> {    const listener = (logged: LoggedPiEvent) => onEvent(logged.event, logged.seq);    this.on("loggedEvent", listener);    this.activePromptCount += 1;    this.lastError = null;    return Effect.tryPromise({      try: () => this.promptSession(message, options),      catch: (error) => error,    }).pipe(      Effect.catch((error) =>        options.restartOnContinuationError !== false && shouldRestartAfterPromptError(error)          ? this.restartPromptEffect(message, options)          : Effect.fail(error),      ),      Effect.catch((error) =>        Effect.sync(() => {          this.lastError = error instanceof Error ? error.message : String(error);        }).pipe(Effect.andThen(Effect.fail(error))),      ),      Effect.ensuring(        Effect.sync(() => {          this.activePromptCount = Math.max(0, this.activePromptCount - 1);          this.off("loggedEvent", listener);        }),      ),    );  }   private promptSession(message: string, options: PiPromptOptions): Promise<void> {    return this.requireSession().prompt(message, {      streamingBehavior: options.streamingBehavior,      images: options.images,      expandPromptTemplates: options.expandPromptTemplates,      source: options.source,      preflightResult: options.preflightResult,    });  }   private restartPromptEffect(    message: string,    options: PiPromptOptions,  ): Effect.Effect<void, unknown> {    return this.ensureStartedEffect(      this.currentModelId,      this.currentCwd,      null,      this.currentStartOptions,    ).pipe(      Effect.andThen(        Effect.tryPromise({          try: () => this.promptSession(message, options),          catch: (error) => error,        }),      ),    );  }   steer(message: string, images: AgentImageInput[] = []): Promise<void> {    return Effect.runPromise(      Effect.tryPromise({        try: () => this.requireSession().steer(message, images),        catch: (error) => error,      }),    );  }   mutateQueuedFollowUp(    message: string,    action: AgentQueueAction,    replacement?: string,    images: AgentImageInput[] = [],  ): Promise<void> {    return Effect.runPromise(      Effect.tryPromise({        try: async () => {          const session = this.requireSession();          this.queueEventBufferDepth += 1;          try {            const cleared = session.clearQueue();

Lines 1–500 of 806.