refactor(coding-agent): add runtime host for session switching closes #2024

This commit is contained in:
Mario Zechner
2026-03-31 13:49:57 +02:00
parent a3bf1eb399
commit d86122cbd3
32 changed files with 1328 additions and 692 deletions

View File

@@ -671,15 +671,21 @@ describe("Context overflow error handling", () => {
}); });
// ============================================================================= // =============================================================================
// llama.cpp server (local) - Skip if not running // llama.cpp server (local) - Skip if not running or not exposing /v1/completions
// ============================================================================= // =============================================================================
let llamaCppRunning = false; let llamaCppRunning = false;
try { if (!process.env.PI_NO_LOCAL_LLM) {
execSync("curl -s --max-time 1 http://localhost:8081/health > /dev/null", { stdio: "ignore" }); try {
llamaCppRunning = true; execSync("curl -s --max-time 1 http://localhost:8081/health > /dev/null", { stdio: "ignore" });
} catch { const probeStatus = execSync(
llamaCppRunning = false; 'curl -s --max-time 1 -o /dev/null -w \'%{http_code}\' -X POST http://localhost:8081/v1/completions -H \'content-type: application/json\' -d \'{"model":"local-model","prompt":"ping","max_tokens":1}\'',
{ encoding: "utf8" },
).trim();
llamaCppRunning = probeStatus !== "404" && probeStatus !== "405" && probeStatus !== "000";
} catch {
llamaCppRunning = false;
}
} }
describe.skipIf(!llamaCppRunning)("llama.cpp (local)", () => { describe.skipIf(!llamaCppRunning)("llama.cpp (local)", () => {

View File

@@ -2,8 +2,77 @@
## [Unreleased] ## [Unreleased]
### Breaking Changes
- Removed extension post-transition events `session_switch` and `session_fork`. Extensions should now use `session_start` and inspect `event.reason`, which is now one of `"startup" | "reload" | "new" | "resume" | "fork"`. For `"new"`, `"resume"`, and `"fork"`, `session_start` also includes `previousSessionFile`. This is better because session replacement now fully reloads extensions, so one post-start hook with explicit reason matches the real lifecycle better than two extra non-cancellable post-transition events.
- Removed session-replacement methods from `AgentSession`. Use `AgentSessionRuntimeHost` for `newSession()`, `switchSession()`, `fork()`, and `importFromJsonl()`. This is better because cross-cwd session replacement rebuilds cwd-bound runtime state and can replace the live `AgentSession` instance entirely. Keeping those operations on a stable runtime host matches the real lifecycle and avoids pretending one `AgentSession` mutates into another.
#### Migration Notes
For existing extensions:
Before:
```ts
pi.on("session_switch", async (event, ctx) => {
if (event.reason === "new") {
resetState();
}
});
pi.on("session_fork", async (_event, ctx) => {
reconstructState(ctx);
});
```
After:
```ts
pi.on("session_start", async (event, ctx) => {
if (event.reason === "new") {
resetState();
}
if (event.reason === "fork" || event.reason === "resume" || event.reason === "startup") {
reconstructState(ctx);
}
});
```
For existing SDK integrations:
Before:
```ts
await session.newSession();
await session.switchSession("/path/to/session.jsonl");
await session.fork("entry-id");
await session.importFromJsonl(jsonl);
```
After:
```ts
const runtime = await createAgentSessionRuntime(bootstrap, {
cwd: process.cwd(),
sessionManager: SessionManager.create(process.cwd()),
});
const runtimeHost = new AgentSessionRuntimeHost(bootstrap, runtime);
await runtimeHost.newSession();
await runtimeHost.switchSession("/path/to/session.jsonl");
await runtimeHost.fork("entry-id");
await runtimeHost.importFromJsonl(jsonl);
const session = runtimeHost.session;
```
After runtime replacement, use `runtimeHost.session` as the new live session and rebind any session-local subscriptions or extension bindings.
### Added ### Added
- Added public SDK runtime-host exports `createAgentSessionRuntime()` and `AgentSessionRuntimeHost` for apps that need runtime-backed session replacement and mode-style session switching
- Added label timestamps to the session tree with a `Shift+T` toggle in `/tree`, smart date formatting, and timestamp preservation through branching ([#2691](https://github.com/badlogic/pi-mono/pull/2691) by [@w-winter](https://github.com/w-winter)) - Added label timestamps to the session tree with a `Shift+T` toggle in `/tree`, smart date formatting, and timestamp preservation through branching ([#2691](https://github.com/badlogic/pi-mono/pull/2691) by [@w-winter](https://github.com/w-winter))
## [0.64.0] - 2026-03-29 ## [0.64.0] - 2026-03-29

View File

@@ -392,15 +392,19 @@ See [docs/packages.md](docs/packages.md).
```typescript ```typescript
import { AuthStorage, createAgentSession, ModelRegistry, SessionManager } from "@mariozechner/pi-coding-agent"; import { AuthStorage, createAgentSession, ModelRegistry, SessionManager } from "@mariozechner/pi-coding-agent";
const authStorage = AuthStorage.create();
const modelRegistry = ModelRegistry.create(authStorage);
const { session } = await createAgentSession({ const { session } = await createAgentSession({
sessionManager: SessionManager.inMemory(), sessionManager: SessionManager.inMemory(),
authStorage: AuthStorage.create(), authStorage,
modelRegistry: ModelRegistry.create(authStorage), modelRegistry,
}); });
await session.prompt("What files are in the current directory?"); await session.prompt("What files are in the current directory?");
``` ```
For advanced multi-session runtime replacement, use `createAgentSessionRuntime()` and `AgentSessionRuntimeHost`.
See [docs/sdk.md](docs/sdk.md) and [examples/sdk/](examples/sdk/). See [docs/sdk.md](docs/sdk.md) and [examples/sdk/](examples/sdk/).
### RPC Mode ### RPC Mode

View File

@@ -37,6 +37,7 @@ See [examples/extensions/](../examples/extensions/) for working implementations.
- [Extension Styles](#extension-styles) - [Extension Styles](#extension-styles)
- [Events](#events) - [Events](#events)
- [Lifecycle Overview](#lifecycle-overview) - [Lifecycle Overview](#lifecycle-overview)
- [Resource Events](#resource-events)
- [Session Events](#session-events) - [Session Events](#session-events)
- [Agent Events](#agent-events) - [Agent Events](#agent-events)
- [Tool Events](#tool-events) - [Tool Events](#tool-events)
@@ -228,7 +229,8 @@ Run `npm install` in the extension directory, then imports from `node_modules/`
pi starts (CLI only) pi starts (CLI only)
├─► session_directory (CLI startup only, no ctx) ├─► session_directory (CLI startup only, no ctx)
─► session_start ─► session_start { reason: "startup" }
└─► resources_discover { reason: "startup" }
user sends prompt ─────────────────────────────────────────┐ user sends prompt ─────────────────────────────────────────┐
@@ -261,11 +263,15 @@ user sends another prompt ◄─────────────────
/new (new session) or /resume (switch session) /new (new session) or /resume (switch session)
├─► session_before_switch (can cancel) ├─► session_before_switch (can cancel)
─► session_switch ─► session_shutdown
├─► session_start { reason: "new" | "resume", previousSessionFile? }
└─► resources_discover { reason: "startup" }
/fork /fork
├─► session_before_fork (can cancel) ├─► session_before_fork (can cancel)
─► session_fork ─► session_shutdown
├─► session_start { reason: "fork", previousSessionFile }
└─► resources_discover { reason: "startup" }
/compact or auto-compaction /compact or auto-compaction
├─► session_before_compact (can cancel or customize) ├─► session_before_compact (can cancel or customize)
@@ -282,6 +288,25 @@ exit (Ctrl+C, Ctrl+D)
└─► session_shutdown └─► session_shutdown
``` ```
### Resource Events
#### resources_discover
Fired after `session_start` so extensions can contribute additional skill, prompt, and theme paths.
The startup path uses `reason: "startup"`. Reload uses `reason: "reload"`.
```typescript
pi.on("resources_discover", async (event, _ctx) => {
// event.cwd - current working directory
// event.reason - "startup" | "reload"
return {
skillPaths: ["/path/to/skills"],
promptPaths: ["/path/to/prompts"],
themePaths: ["/path/to/themes"],
};
});
```
### Session Events ### Session Events
See [session.md](session.md) for session storage internals and the SessionManager API. See [session.md](session.md) for session storage internals and the SessionManager API.
@@ -309,17 +334,19 @@ pi.on("session_directory", async (event) => {
#### session_start #### session_start
Fired on initial session load. Fired when a session is started, loaded, or reloaded.
```typescript ```typescript
pi.on("session_start", async (_event, ctx) => { pi.on("session_start", async (event, ctx) => {
// event.reason - "startup" | "reload" | "new" | "resume" | "fork"
// event.previousSessionFile - present for "new", "resume", and "fork"
ctx.ui.notify(`Session: ${ctx.sessionManager.getSessionFile() ?? "ephemeral"}`, "info"); ctx.ui.notify(`Session: ${ctx.sessionManager.getSessionFile() ?? "ephemeral"}`, "info");
}); });
``` ```
#### session_before_switch / session_switch #### session_before_switch
Fired when starting a new session (`/new`) or switching sessions (`/resume`). Fired before starting a new session (`/new`) or switching sessions (`/resume`).
```typescript ```typescript
pi.on("session_before_switch", async (event, ctx) => { pi.on("session_before_switch", async (event, ctx) => {
@@ -331,14 +358,12 @@ pi.on("session_before_switch", async (event, ctx) => {
if (!ok) return { cancel: true }; if (!ok) return { cancel: true };
} }
}); });
pi.on("session_switch", async (event, ctx) => {
// event.reason - "new" or "resume"
// event.previousSessionFile - session we came from
});
``` ```
#### session_before_fork / session_fork After a successful switch or new-session action, pi emits `session_shutdown` for the old extension instance, reloads and rebinds extensions for the new session, then emits `session_start` with `reason: "new" | "resume"` and `previousSessionFile`.
Do cleanup work in `session_shutdown`, then reestablish any in-memory state in `session_start`.
#### session_before_fork
Fired when forking via `/fork`. Fired when forking via `/fork`.
@@ -349,12 +374,11 @@ pi.on("session_before_fork", async (event, ctx) => {
// OR // OR
return { skipConversationRestore: true }; // Fork but don't rewind messages return { skipConversationRestore: true }; // Fork but don't rewind messages
}); });
pi.on("session_fork", async (event, ctx) => {
// event.previousSessionFile - previous session file
});
``` ```
After a successful fork, pi emits `session_shutdown` for the old extension instance, reloads and rebinds extensions for the new session, then emits `session_start` with `reason: "fork"` and `previousSessionFile`.
Do cleanup work in `session_shutdown`, then reestablish any in-memory state in `session_start`.
#### session_before_compact / session_compact #### session_before_compact / session_compact
Fired on compaction. See [compaction.md](compaction.md) for details. Fired on compaction. See [compaction.md](compaction.md) for details.
@@ -936,7 +960,7 @@ pi.registerCommand("reload-runtime", {
Important behavior: Important behavior:
- `await ctx.reload()` emits `session_shutdown` for the current extension runtime - `await ctx.reload()` emits `session_shutdown` for the current extension runtime
- It then reloads resources and emits `session_start` (and `resources_discover` with reason `"reload"`) for the new runtime - It then reloads resources and emits `session_start` with `reason: "reload"` and `resources_discover` with reason `"reload"`
- The currently running command handler still continues in the old call frame - The currently running command handler still continues in the old call frame
- Code after `await ctx.reload()` still runs from the pre-reload version - Code after `await ctx.reload()` still runs from the pre-reload version
- Code after `await ctx.reload()` must not assume old in-memory extension state is still valid - Code after `await ctx.reload()` must not assume old in-memory extension state is still valid
@@ -2187,7 +2211,7 @@ All examples in [examples/extensions/](../examples/extensions/).
| **Compaction & Sessions** ||| | **Compaction & Sessions** |||
| `custom-compaction.ts` | Custom compaction summary | `on("session_before_compact")` | | `custom-compaction.ts` | Custom compaction summary | `on("session_before_compact")` |
| `trigger-compact.ts` | Trigger compaction manually | `compact()` | | `trigger-compact.ts` | Trigger compaction manually | `compact()` |
| `git-checkpoint.ts` | Git stash on turns | `on("turn_end")`, `on("session_fork")`, `exec` | | `git-checkpoint.ts` | Git stash on turns | `on("turn_start")`, `on("session_before_fork")`, `exec` |
| `auto-commit-on-exit.ts` | Commit on shutdown | `on("session_shutdown")`, `exec` | | `auto-commit-on-exit.ts` | Commit on shutdown | `on("session_shutdown")`, `exec` |
| **UI Components** ||| | **UI Components** |||
| `status-line.ts` | Footer status indicator | `setStatus`, session events | | `status-line.ts` | Footer status indicator | `setStatus`, session events |

View File

@@ -49,7 +49,7 @@ The SDK is included in the main package. No separate installation needed.
### createAgentSession() ### createAgentSession()
The main factory function. Creates an `AgentSession` with configurable options. The main factory function for a single `AgentSession`.
`createAgentSession()` uses a `ResourceLoader` to supply extensions, skills, prompt templates, themes, and context files. If you do not provide one, it uses `DefaultResourceLoader` with standard discovery. `createAgentSession()` uses a `ResourceLoader` to supply extensions, skills, prompt templates, themes, and context files. If you do not provide one, it uses `DefaultResourceLoader` with standard discovery.
@@ -69,61 +69,102 @@ const { session } = await createAgentSession({
### AgentSession ### AgentSession
The session manages the agent lifecycle, message history, and event streaming. The session manages agent lifecycle, message history, model state, compaction, and event streaming.
```typescript ```typescript
interface AgentSession { interface AgentSession {
// Send a prompt and wait for completion // Send a prompt and wait for completion
// If streaming, requires streamingBehavior option to queue the message
prompt(text: string, options?: PromptOptions): Promise<void>; prompt(text: string, options?: PromptOptions): Promise<void>;
// Queue messages during streaming // Queue messages during streaming
steer(text: string): Promise<void>; // Queue for delivery after the current assistant turn finishes its tool calls steer(text: string): Promise<void>;
followUp(text: string): Promise<void>; // Wait: delivered only when agent finishes followUp(text: string): Promise<void>;
// Subscribe to events (returns unsubscribe function) // Subscribe to events (returns unsubscribe function)
subscribe(listener: (event: AgentSessionEvent) => void): () => void; subscribe(listener: (event: AgentSessionEvent) => void): () => void;
// Session info // Session info
sessionFile: string | undefined; // undefined for in-memory sessionFile: string | undefined;
sessionId: string; sessionId: string;
// Model control // Model control
setModel(model: Model): Promise<void>; setModel(model: Model): Promise<void>;
setThinkingLevel(level: ThinkingLevel): void; setThinkingLevel(level: ThinkingLevel): void;
cycleModel(): Promise<ModelCycleResult | undefined>; cycleModel(): Promise<ModelCycleResult | undefined>;
cycleThinkingLevel(): ThinkingLevel | undefined; cycleThinkingLevel(): ThinkingLevel | undefined;
// State access // State access
agent: Agent; agent: Agent;
model: Model | undefined; model: Model | undefined;
thinkingLevel: ThinkingLevel; thinkingLevel: ThinkingLevel;
messages: AgentMessage[]; messages: AgentMessage[];
isStreaming: boolean; isStreaming: boolean;
// Session management // In-place tree navigation within the current session file
newSession(options?: { parentSession?: string }): Promise<boolean>; // Returns false if cancelled by hook navigateTree(targetId: string, options?: { summarize?: boolean; customInstructions?: string; replaceInstructions?: boolean; label?: string }): Promise<{ editorText?: string; cancelled: boolean }>;
switchSession(sessionPath: string): Promise<boolean>;
// Forking
fork(entryId: string): Promise<{ selectedText: string; cancelled: boolean }>; // Creates new session file
navigateTree(targetId: string, options?: { summarize?: boolean; customInstructions?: string; replaceInstructions?: boolean; label?: string }): Promise<{ editorText?: string; cancelled: boolean }>; // In-place navigation
// Hook message injection
sendHookMessage(message: HookMessage, triggerTurn?: boolean): Promise<void>;
// Compaction // Compaction
compact(customInstructions?: string): Promise<CompactionResult>; compact(customInstructions?: string): Promise<CompactionResult>;
abortCompaction(): void; abortCompaction(): void;
// Abort current operation // Abort current operation
abort(): Promise<void>; abort(): Promise<void>;
// Cleanup // Cleanup
dispose(): void; dispose(): void;
} }
``` ```
Session replacement APIs such as new-session, resume, fork, and import live on `AgentSessionRuntimeHost`, not on `AgentSession`.
### createAgentSessionRuntime() and AgentSessionRuntimeHost
Use the runtime API when you need to replace the active session and rebuild cwd-bound runtime state.
This is the same layer used by the built-in interactive, print, and RPC modes.
```typescript
import {
AgentSessionRuntimeHost,
createAgentSessionRuntime,
SessionManager,
} from "@mariozechner/pi-coding-agent";
const bootstrap = {
// Optional: authStorage, model, thinkingLevel, tools, customTools, resourceLoader
};
const runtime = await createAgentSessionRuntime(bootstrap, {
cwd: process.cwd(),
sessionManager: SessionManager.create(process.cwd()),
});
const runtimeHost = new AgentSessionRuntimeHost(bootstrap, runtime);
```
`createAgentSessionRuntime()` returns an internal runtime bundle. `AgentSessionRuntimeHost` owns replacement of that bundle across:
- `newSession()`
- `switchSession()`
- `fork()`
- `importFromJsonl()`
Important behavior:
- `runtimeHost.session` changes after those operations
- event subscriptions are attached to a specific `AgentSession`, so re-subscribe after replacement
- if you use extensions, call `runtimeHost.session.bindExtensions(...)` again for the new session
```typescript
let session = runtimeHost.session;
let unsubscribe = session.subscribe(() => {});
await runtimeHost.newSession();
unsubscribe();
session = runtimeHost.session;
unsubscribe = session.subscribe(() => {});
```
### Prompting and Message Queueing ### Prompting and Message Queueing
The `prompt()` method handles prompt templates, extension commands, and message sending: The `prompt()` method handles prompt templates, extension commands, and message sending:
@@ -602,7 +643,12 @@ const { session } = await createAgentSession({ resourceLoader: loader });
Sessions use a tree structure with `id`/`parentId` linking, enabling in-place branching. Sessions use a tree structure with `id`/`parentId` linking, enabling in-place branching.
```typescript ```typescript
import { createAgentSession, SessionManager } from "@mariozechner/pi-coding-agent"; import {
AgentSessionRuntimeHost,
createAgentSession,
createAgentSessionRuntime,
SessionManager,
} from "@mariozechner/pi-coding-agent";
// In-memory (no persistence) // In-memory (no persistence)
const { session } = await createAgentSession({ const { session } = await createAgentSession({
@@ -610,12 +656,12 @@ const { session } = await createAgentSession({
}); });
// New persistent session // New persistent session
const { session } = await createAgentSession({ const { session: persisted } = await createAgentSession({
sessionManager: SessionManager.create(process.cwd()), sessionManager: SessionManager.create(process.cwd()),
}); });
// Continue most recent // Continue most recent
const { session, modelFallbackMessage } = await createAgentSession({ const { session: continued, modelFallbackMessage } = await createAgentSession({
sessionManager: SessionManager.continueRecent(process.cwd()), sessionManager: SessionManager.continueRecent(process.cwd()),
}); });
if (modelFallbackMessage) { if (modelFallbackMessage) {
@@ -623,26 +669,30 @@ if (modelFallbackMessage) {
} }
// Open specific file // Open specific file
const { session } = await createAgentSession({ const { session: opened } = await createAgentSession({
sessionManager: SessionManager.open("/path/to/session.jsonl"), sessionManager: SessionManager.open("/path/to/session.jsonl"),
}); });
// List available sessions (async with optional progress callback) // List sessions
const sessions = await SessionManager.list(process.cwd()); const currentProjectSessions = await SessionManager.list(process.cwd());
for (const info of sessions) { const allSessions = await SessionManager.listAll(process.cwd());
console.log(`${info.id}: ${info.firstMessage} (${info.messageCount} messages, cwd: ${info.cwd})`);
}
// List all sessions across all projects // Session replacement API for /new, /resume, /fork, and import flows.
const allSessions = await SessionManager.listAll((loaded, total) => { const bootstrap = {};
console.log(`Loading ${loaded}/${total}...`); const runtime = await createAgentSessionRuntime(bootstrap, {
cwd: process.cwd(),
sessionManager: SessionManager.create(process.cwd()),
}); });
const runtimeHost = new AgentSessionRuntimeHost(bootstrap, runtime);
// Custom session directory (no cwd encoding) // Replace the active session with a fresh one
const customDir = "/path/to/my-sessions"; await runtimeHost.newSession();
const { session } = await createAgentSession({
sessionManager: SessionManager.create(process.cwd(), customDir), // Replace the active session with another saved session
}); await runtimeHost.switchSession("/path/to/session.jsonl");
// Replace the active session with a fork from a specific entry
await runtimeHost.fork("entry-id");
``` ```
**SessionManager tree API:** **SessionManager tree API:**
@@ -650,6 +700,10 @@ const { session } = await createAgentSession({
```typescript ```typescript
const sm = SessionManager.open("/path/to/session.jsonl"); const sm = SessionManager.open("/path/to/session.jsonl");
// Session listing
const currentProjectSessions = await SessionManager.list(process.cwd());
const allSessions = await SessionManager.listAll(process.cwd());
// Tree traversal // Tree traversal
const entries = sm.getEntries(); // All entries (excludes header) const entries = sm.getEntries(); // All entries (excludes header)
const tree = sm.getTree(); // Full tree structure const tree = sm.getTree(); // Full tree structure
@@ -859,20 +913,29 @@ The SDK exports run mode utilities for building custom interfaces on top of `cre
Full TUI interactive mode with editor, chat history, and all built-in commands: Full TUI interactive mode with editor, chat history, and all built-in commands:
```typescript ```typescript
import { createAgentSession, InteractiveMode } from "@mariozechner/pi-coding-agent"; import {
AgentSessionRuntimeHost,
createAgentSessionRuntime,
InteractiveMode,
SessionManager,
} from "@mariozechner/pi-coding-agent";
const { session } = await createAgentSession({ /* ... */ }); const bootstrap = {};
const runtime = await createAgentSessionRuntime(bootstrap, {
cwd: process.cwd(),
sessionManager: SessionManager.create(process.cwd()),
});
const runtimeHost = new AgentSessionRuntimeHost(bootstrap, runtime);
const mode = new InteractiveMode(session, { const mode = new InteractiveMode(runtimeHost, {
// All optional migratedProviders: [],
migratedProviders: [], // Show migration warnings modelFallbackMessage: undefined,
modelFallbackMessage: undefined, // Show model restore warning initialMessage: "Hello",
initialMessage: "Hello", // Send on startup initialImages: [],
initialImages: [], // Images with initial message initialMessages: [],
initialMessages: [], // Additional startup prompts
}); });
await mode.run(); // Blocks until exit await mode.run();
``` ```
### runPrintMode ### runPrintMode
@@ -880,15 +943,25 @@ await mode.run(); // Blocks until exit
Single-shot mode: send prompts, output result, exit: Single-shot mode: send prompts, output result, exit:
```typescript ```typescript
import { createAgentSession, runPrintMode } from "@mariozechner/pi-coding-agent"; import {
AgentSessionRuntimeHost,
createAgentSessionRuntime,
runPrintMode,
SessionManager,
} from "@mariozechner/pi-coding-agent";
const { session } = await createAgentSession({ /* ... */ }); const bootstrap = {};
const runtime = await createAgentSessionRuntime(bootstrap, {
cwd: process.cwd(),
sessionManager: SessionManager.create(process.cwd()),
});
const runtimeHost = new AgentSessionRuntimeHost(bootstrap, runtime);
await runPrintMode(session, { await runPrintMode(runtimeHost, {
mode: "text", // "text" for final response, "json" for all events mode: "text",
initialMessage: "Hello", // First message (can include @file content) initialMessage: "Hello",
initialImages: [], // Images with initial message initialImages: [],
messages: ["Follow up"], // Additional prompts messages: ["Follow up"],
}); });
``` ```
@@ -897,11 +970,21 @@ await runPrintMode(session, {
JSON-RPC mode for subprocess integration: JSON-RPC mode for subprocess integration:
```typescript ```typescript
import { createAgentSession, runRpcMode } from "@mariozechner/pi-coding-agent"; import {
AgentSessionRuntimeHost,
createAgentSessionRuntime,
runRpcMode,
SessionManager,
} from "@mariozechner/pi-coding-agent";
const { session } = await createAgentSession({ /* ... */ }); const bootstrap = {};
const runtime = await createAgentSessionRuntime(bootstrap, {
cwd: process.cwd(),
sessionManager: SessionManager.create(process.cwd()),
});
const runtimeHost = new AgentSessionRuntimeHost(bootstrap, runtime);
await runRpcMode(session); // Reads JSON commands from stdin, writes to stdout await runRpcMode(runtimeHost);
``` ```
See [RPC documentation](rpc.md) for the JSON protocol. See [RPC documentation](rpc.md) for the JSON protocol.
@@ -934,6 +1017,8 @@ The main entry point exports:
```typescript ```typescript
// Factory // Factory
createAgentSession createAgentSession
createAgentSessionRuntime
AgentSessionRuntimeHost
// Auth and Models // Auth and Models
AuthStorage AuthStorage

View File

@@ -13,7 +13,7 @@ Sessions are stored as trees where each entry has an `id` and `parentId`. The "l
| View | Flat list of user messages | Full tree structure | | View | Flat list of user messages | Full tree structure |
| Action | Extracts path to **new session file** | Changes leaf in **same session** | | Action | Extracts path to **new session file** | Changes leaf in **same session** |
| Summary | Never | Optional (user prompted) | | Summary | Never | Optional (user prompted) |
| Events | `session_before_fork` / `session_fork` | `session_before_tree` / `session_tree` | | Events | `session_before_fork` / `session_start` (`reason: "fork"`) | `session_before_tree` / `session_tree` |
## Tree UI ## Tree UI

View File

@@ -33,10 +33,6 @@ export default function (pi: ExtensionAPI) {
applyLabel(ctx); applyLabel(ctx);
}); });
pi.on("session_switch", async (_event, ctx) => {
applyLabel(ctx);
});
pi.registerCommand("thinking-label", { pi.registerCommand("thinking-label", {
description: "Set the hidden thinking label. Use without args to reset.", description: "Set the hidden thinking label. Use without args to reset.",
handler: async (args, ctx) => { handler: async (args, ctx) => {

View File

@@ -13,7 +13,7 @@
* - notify() - after each dialog completes * - notify() - after each dialog completes
* - setStatus() - on turn_start/turn_end * - setStatus() - on turn_start/turn_end
* - setWidget() - on session_start * - setWidget() - on session_start
* - setTitle() - on session_start and session_switch * - setTitle() - on session_start
* - setEditorText() - via /rpc-prefill command * - setEditorText() - via /rpc-prefill command
*/ */
@@ -24,18 +24,12 @@ export default function (pi: ExtensionAPI) {
// -- setTitle, setWidget, setStatus on session lifecycle -- // -- setTitle, setWidget, setStatus on session lifecycle --
pi.on("session_start", async (_event, ctx) => { pi.on("session_start", async (event, ctx) => {
ctx.ui.setTitle("pi RPC Demo"); ctx.ui.setTitle(event.reason === "new" ? "pi RPC Demo (new session)" : "pi RPC Demo");
ctx.ui.setWidget("rpc-demo", ["--- RPC Extension UI Demo ---", "Loaded and ready."]); ctx.ui.setWidget("rpc-demo", ["--- RPC Extension UI Demo ---", "Loaded and ready."]);
ctx.ui.setStatus("rpc-demo", `Turns: ${turnCount}`); ctx.ui.setStatus("rpc-demo", `Turns: ${turnCount}`);
}); });
pi.on("session_switch", async (_event, ctx) => {
turnCount = 0;
ctx.ui.setTitle("pi RPC Demo (new session)");
ctx.ui.setStatus("rpc-demo", `Turns: ${turnCount}`);
});
// -- setStatus on turn lifecycle -- // -- setStatus on turn lifecycle --
pi.on("turn_start", async (_event, ctx) => { pi.on("turn_start", async (_event, ctx) => {

View File

@@ -29,12 +29,4 @@ export default function (pi: ExtensionAPI) {
const text = theme.fg("dim", ` Turn ${turnCount} complete`); const text = theme.fg("dim", ` Turn ${turnCount} complete`);
ctx.ui.setStatus("status-demo", check + text); ctx.ui.setStatus("status-demo", check + text);
}); });
pi.on("session_switch", async (event, ctx) => {
if (event.reason === "new") {
turnCount = 0;
const theme = ctx.ui.theme;
ctx.ui.setStatus("status-demo", theme.fg("dim", "Ready"));
}
});
} }

View File

@@ -130,8 +130,6 @@ export default function (pi: ExtensionAPI) {
// Reconstruct state on session events // Reconstruct state on session events
pi.on("session_start", async (_event, ctx) => reconstructState(ctx)); pi.on("session_start", async (_event, ctx) => reconstructState(ctx));
pi.on("session_switch", async (_event, ctx) => reconstructState(ctx));
pi.on("session_fork", async (_event, ctx) => reconstructState(ctx));
pi.on("session_tree", async (_event, ctx) => reconstructState(ctx)); pi.on("session_tree", async (_event, ctx) => reconstructState(ctx));
// Register the todo tool for the LLM // Register the todo tool for the LLM

View File

@@ -138,9 +138,4 @@ export default function toolsExtension(pi: ExtensionAPI) {
pi.on("session_tree", async (_event, ctx) => { pi.on("session_tree", async (_event, ctx) => {
restoreFromBranch(ctx); restoreFromBranch(ctx);
}); });
// Restore state after forking
pi.on("session_fork", async (_event, ctx) => {
restoreFromBranch(ctx);
});
} }

View File

@@ -1,17 +1,9 @@
import type { ExtensionAPI, ExtensionContext } from "@mariozechner/pi-coding-agent"; import type { ExtensionAPI } from "@mariozechner/pi-coding-agent";
const applyWidgets = (ctx: ExtensionContext) => {
if (!ctx.hasUI) return;
ctx.ui.setWidget("widget-above", ["Above editor widget"]);
ctx.ui.setWidget("widget-below", ["Below editor widget"], { placement: "belowEditor" });
};
export default function widgetPlacementExtension(pi: ExtensionAPI) { export default function widgetPlacementExtension(pi: ExtensionAPI) {
pi.on("session_start", (_event, ctx) => { pi.on("session_start", (_event, ctx) => {
applyWidgets(ctx); if (!ctx.hasUI) return;
}); ctx.ui.setWidget("widget-above", ["Above editor widget"]);
ctx.ui.setWidget("widget-below", ["Below editor widget"], { placement: "belowEditor" });
pi.on("session_switch", (_event, ctx) => {
applyWidgets(ctx);
}); });
} }

View File

@@ -0,0 +1,49 @@
/**
* Session Runtime Host
*
* Use the runtime host when you need to replace the active AgentSession,
* for example for new-session, resume, fork, or import flows.
*
* The important pattern is: after the host replaces the runtime, rebind any
* session-local subscriptions and extension bindings to `runtimeHost.session`.
*/
import { AgentSessionRuntimeHost, createAgentSessionRuntime, SessionManager } from "@mariozechner/pi-coding-agent";
const bootstrap = {};
const runtime = await createAgentSessionRuntime(bootstrap, {
cwd: process.cwd(),
sessionManager: SessionManager.create(process.cwd()),
});
const runtimeHost = new AgentSessionRuntimeHost(bootstrap, runtime);
let unsubscribe: (() => void) | undefined;
async function bindSession() {
unsubscribe?.();
const session = runtimeHost.session;
await session.bindExtensions({});
unsubscribe = session.subscribe((event) => {
if (event.type === "queue_update") {
console.log("Queued:", event.steering.length + event.followUp.length);
}
});
return session;
}
let session = await bindSession();
const originalSessionFile = session.sessionFile;
console.log("Initial session:", originalSessionFile);
await runtimeHost.newSession();
session = await bindSession();
console.log("After newSession():", session.sessionFile);
if (originalSessionFile) {
await runtimeHost.switchSession(originalSessionFile);
session = await bindSession();
console.log("After switchSession():", session.sessionFile);
}
unsubscribe?.();
await runtimeHost.dispose();

View File

@@ -1,6 +1,6 @@
# SDK Examples # SDK Examples
Programmatic usage of pi-coding-agent via `createAgentSession()`. Programmatic usage of pi-coding-agent via `createAgentSession()` and `createAgentSessionRuntime()`.
## Examples ## Examples
@@ -18,6 +18,7 @@ Programmatic usage of pi-coding-agent via `createAgentSession()`.
| `10-settings.ts` | Override compaction, retry, terminal settings | | `10-settings.ts` | Override compaction, retry, terminal settings |
| `11-sessions.ts` | In-memory, persistent, continue, list sessions | | `11-sessions.ts` | In-memory, persistent, continue, list sessions |
| `12-full-control.ts` | Replace everything, no discovery | | `12-full-control.ts` | Replace everything, no discovery |
| `13-session-runtime.ts` | Manage runtime-backed session replacement |
## Running ## Running

View File

@@ -0,0 +1,358 @@
import { copyFileSync, existsSync, mkdirSync } from "node:fs";
import { basename, join, resolve } from "node:path";
import type { ThinkingLevel } from "@mariozechner/pi-agent-core";
import type { Model } from "@mariozechner/pi-ai";
import { getAgentDir } from "../config.js";
import { AuthStorage } from "./auth-storage.js";
import type { SessionStartEvent, ToolDefinition } from "./extensions/index.js";
import { emitSessionShutdownEvent } from "./extensions/runner.js";
import { ModelRegistry } from "./model-registry.js";
import { DefaultResourceLoader, type DefaultResourceLoaderOptions, type ResourceLoader } from "./resource-loader.js";
import { type CreateAgentSessionResult, createAgentSession } from "./sdk.js";
import { SessionManager } from "./session-manager.js";
import { SettingsManager } from "./settings-manager.js";
import type { Tool } from "./tools/index.js";
/**
* Stable bootstrap inputs reused whenever the active runtime is replaced.
*
* Use this for process-level wiring that should survive `/new`, `/resume`,
* `/fork`, and import flows. Session-local state belongs in the session file
* or in settings, not here.
*/
export interface AgentSessionRuntimeBootstrap {
/** Agent directory used for auth, models, settings, sessions, and resource discovery. */
agentDir?: string;
/** Optional shared auth storage. If omitted, file-backed storage under agentDir is used. */
authStorage?: AuthStorage;
/** Initial model for the first created session runtime. */
model?: Model<any>;
/** Initial thinking level for the first created session runtime. */
thinkingLevel?: ThinkingLevel;
/** Optional scoped model list for model cycling and selection. */
scopedModels?: Array<{ model: Model<any>; thinkingLevel?: ThinkingLevel }>;
/** Built-in tool set override. */
tools?: Tool[];
/** Additional custom tools registered directly through the SDK. */
customTools?: ToolDefinition[];
/**
* Resource loader input used for each created runtime.
*
* Pass either a factory that creates a fully custom ResourceLoader for the
* target cwd, or DefaultResourceLoader options without cwd/agentDir/
* settingsManager, which are supplied by the runtime.
*/
resourceLoader?:
| ((cwd: string, agentDir: string) => Promise<ResourceLoader>)
| Omit<DefaultResourceLoaderOptions, "cwd" | "agentDir" | "settingsManager">;
}
/** Options for creating one concrete runtime instance. */
export interface CreateAgentSessionRuntimeOptions {
/** Working directory for this runtime instance. */
cwd: string;
/** Optional preselected session manager. If omitted, normal session resolution applies. */
sessionManager?: SessionManager;
/** Optional session_start metadata to emit when the runtime binds extensions. */
sessionStartEvent?: SessionStartEvent;
}
type AgentSessionRuntime = CreateAgentSessionResult & {
cwd: string;
agentDir: string;
authStorage: AuthStorage;
modelRegistry: ModelRegistry;
settingsManager: SettingsManager;
resourceLoader: ResourceLoader;
sessionManager: SessionManager;
};
/**
* Create one runtime instance containing an AgentSession plus the cwd-bound
* services it depends on.
*
* Most SDK callers should keep the returned value wrapped in an
* AgentSessionRuntimeHost instead of holding it directly. The host owns
* replacing the runtime when switching sessions across files or working
* directories.
*/
export async function createAgentSessionRuntime(
bootstrap: AgentSessionRuntimeBootstrap,
options: CreateAgentSessionRuntimeOptions,
): Promise<AgentSessionRuntime> {
const cwd = options.cwd;
const agentDir = bootstrap.agentDir ?? getAgentDir();
const authStorage = bootstrap.authStorage ?? AuthStorage.create(join(agentDir, "auth.json"));
const settingsManager = SettingsManager.create(cwd, agentDir);
const modelRegistry = ModelRegistry.create(authStorage, join(agentDir, "models.json"));
const resourceLoader =
typeof bootstrap.resourceLoader === "function"
? await bootstrap.resourceLoader(cwd, agentDir)
: new DefaultResourceLoader({
...(bootstrap.resourceLoader ?? {}),
cwd,
agentDir,
settingsManager,
});
await resourceLoader.reload();
const extensionsResult = resourceLoader.getExtensions();
for (const { name, config } of extensionsResult.runtime.pendingProviderRegistrations) {
modelRegistry.registerProvider(name, config);
}
extensionsResult.runtime.pendingProviderRegistrations = [];
const created = await createAgentSession({
cwd,
agentDir,
authStorage,
modelRegistry,
settingsManager,
resourceLoader,
sessionManager: options.sessionManager,
model: bootstrap.model,
thinkingLevel: bootstrap.thinkingLevel,
scopedModels: bootstrap.scopedModels,
tools: bootstrap.tools,
customTools: bootstrap.customTools,
sessionStartEvent: options.sessionStartEvent,
});
return {
...created,
cwd,
agentDir,
authStorage,
modelRegistry,
settingsManager,
resourceLoader,
sessionManager: created.session.sessionManager,
};
}
function extractUserMessageText(content: string | Array<{ type: string; text?: string }>): string {
if (typeof content === "string") {
return content;
}
return content
.filter((part): part is { type: "text"; text: string } => part.type === "text" && typeof part.text === "string")
.map((part) => part.text)
.join("");
}
/**
* Stable wrapper around a replaceable AgentSession runtime.
*
* Use this when your application needs `/new`, `/resume`, `/fork`, or import
* behavior. After replacement, read `session` again and rebind any
* session-local subscriptions or extension bindings.
*/
export class AgentSessionRuntimeHost {
constructor(
private readonly bootstrap: AgentSessionRuntimeBootstrap,
private runtime: AgentSessionRuntime,
) {}
/** The currently active session instance. Re-read this after runtime replacement. */
get session() {
return this.runtime.session;
}
private async emitBeforeSwitch(
reason: "new" | "resume",
targetSessionFile?: string,
): Promise<{ cancelled: boolean }> {
const runner = this.runtime.session.extensionRunner;
if (!runner?.hasHandlers("session_before_switch")) {
return { cancelled: false };
}
const result = await runner.emit({
type: "session_before_switch",
reason,
targetSessionFile,
});
return { cancelled: result?.cancel === true };
}
private async emitBeforeFork(entryId: string): Promise<{ cancelled: boolean }> {
const runner = this.runtime.session.extensionRunner;
if (!runner?.hasHandlers("session_before_fork")) {
return { cancelled: false };
}
const result = await runner.emit({
type: "session_before_fork",
entryId,
});
return { cancelled: result?.cancel === true };
}
private async replace(options: CreateAgentSessionRuntimeOptions): Promise<void> {
const nextRuntime = await createAgentSessionRuntime(this.bootstrap, options);
await emitSessionShutdownEvent(this.runtime.session.extensionRunner);
this.runtime.session.dispose();
if (process.cwd() !== nextRuntime.cwd) {
process.chdir(nextRuntime.cwd);
}
this.runtime = nextRuntime;
}
/**
* Replace the active runtime with one opened from an existing session file.
*
* Emits `session_before_switch` before replacement and returns
* `{ cancelled: true }` if an extension vetoes the switch.
*/
async switchSession(sessionPath: string): Promise<{ cancelled: boolean }> {
const beforeResult = await this.emitBeforeSwitch("resume", sessionPath);
if (beforeResult.cancelled) {
return beforeResult;
}
const previousSessionFile = this.runtime.session.sessionFile;
const sessionManager = SessionManager.open(sessionPath);
await this.replace({
cwd: sessionManager.getCwd(),
sessionManager,
sessionStartEvent: { type: "session_start", reason: "resume", previousSessionFile },
});
return { cancelled: false };
}
/**
* Replace the active runtime with a fresh session in the current cwd.
*
* `setup` runs after replacement against the new session manager, which lets
* callers seed entries before normal use begins.
*/
async newSession(options?: {
/** Optional parent session path recorded in the new session header. */
parentSession?: string;
/** Optional callback for seeding the new session manager after replacement. */
setup?: (sessionManager: SessionManager) => Promise<void>;
}): Promise<{ cancelled: boolean }> {
const beforeResult = await this.emitBeforeSwitch("new");
if (beforeResult.cancelled) {
return beforeResult;
}
const previousSessionFile = this.runtime.session.sessionFile;
const sessionDir = this.runtime.sessionManager.getSessionDir();
const sessionManager = SessionManager.create(this.runtime.cwd, sessionDir);
if (options?.parentSession) {
sessionManager.newSession({ parentSession: options.parentSession });
}
await this.replace({
cwd: this.runtime.cwd,
sessionManager,
sessionStartEvent: { type: "session_start", reason: "new", previousSessionFile },
});
if (options?.setup) {
await options.setup(this.runtime.sessionManager);
this.runtime.session.agent.state.messages = this.runtime.sessionManager.buildSessionContext().messages;
}
return { cancelled: false };
}
/**
* Replace the active runtime with a fork rooted at the given user-message
* entry.
*
* Returns the selected user text so UIs can restore it into the editor after
* the fork completes.
*/
async fork(entryId: string): Promise<{ cancelled: boolean; selectedText?: string }> {
const beforeResult = await this.emitBeforeFork(entryId);
if (beforeResult.cancelled) {
return { cancelled: true };
}
const selectedEntry = this.runtime.sessionManager.getEntry(entryId);
if (!selectedEntry || selectedEntry.type !== "message" || selectedEntry.message.role !== "user") {
throw new Error("Invalid entry ID for forking");
}
const previousSessionFile = this.runtime.session.sessionFile;
const selectedText = extractUserMessageText(selectedEntry.message.content);
if (this.runtime.sessionManager.isPersisted()) {
const currentSessionFile = this.runtime.session.sessionFile!;
const sessionDir = this.runtime.sessionManager.getSessionDir();
if (!selectedEntry.parentId) {
const sessionManager = SessionManager.create(this.runtime.cwd, sessionDir);
sessionManager.newSession({ parentSession: currentSessionFile });
await this.replace({
cwd: this.runtime.cwd,
sessionManager,
sessionStartEvent: { type: "session_start", reason: "fork", previousSessionFile },
});
return { cancelled: false, selectedText };
}
const sourceManager = SessionManager.open(currentSessionFile, sessionDir);
const forkedSessionPath = sourceManager.createBranchedSession(selectedEntry.parentId)!;
const sessionManager = SessionManager.open(forkedSessionPath, sessionDir);
await this.replace({
cwd: sessionManager.getCwd(),
sessionManager,
sessionStartEvent: { type: "session_start", reason: "fork", previousSessionFile },
});
return { cancelled: false, selectedText };
}
const sessionManager = this.runtime.sessionManager;
if (!selectedEntry.parentId) {
sessionManager.newSession({ parentSession: this.runtime.session.sessionFile });
} else {
sessionManager.createBranchedSession(selectedEntry.parentId);
}
await this.replace({
cwd: this.runtime.cwd,
sessionManager,
sessionStartEvent: { type: "session_start", reason: "fork", previousSessionFile },
});
return { cancelled: false, selectedText };
}
/**
* Import a JSONL session file into the current session directory and replace
* the active runtime with the imported session.
*/
async importFromJsonl(inputPath: string): Promise<{ cancelled: boolean }> {
const resolvedPath = resolve(inputPath);
if (!existsSync(resolvedPath)) {
throw new Error(`File not found: ${resolvedPath}`);
}
const sessionDir = this.runtime.sessionManager.getSessionDir();
if (!existsSync(sessionDir)) {
mkdirSync(sessionDir, { recursive: true });
}
const destinationPath = join(sessionDir, basename(resolvedPath));
const beforeResult = await this.emitBeforeSwitch("resume", destinationPath);
if (beforeResult.cancelled) {
return beforeResult;
}
const previousSessionFile = this.runtime.session.sessionFile;
if (resolve(destinationPath) !== resolvedPath) {
copyFileSync(resolvedPath, destinationPath);
}
const sessionManager = SessionManager.open(destinationPath, sessionDir);
await this.replace({
cwd: sessionManager.getCwd(),
sessionManager,
sessionStartEvent: { type: "session_start", reason: "resume", previousSessionFile },
});
return { cancelled: false };
}
/** Emit session shutdown for the active runtime and dispose it permanently. */
async dispose(): Promise<void> {
await emitSessionShutdownEvent(this.runtime.session.extensionRunner);
this.runtime.session.dispose();
}
}

View File

@@ -13,7 +13,7 @@
* Modes use this class and add their own I/O layer on top. * Modes use this class and add their own I/O layer on top.
*/ */
import { copyFileSync, existsSync, mkdirSync, readFileSync, writeFileSync } from "node:fs"; import { existsSync, mkdirSync, readFileSync, writeFileSync } from "node:fs";
import { basename, dirname, join, resolve } from "node:path"; import { basename, dirname, join, resolve } from "node:path";
import type { import type {
Agent, Agent,
@@ -29,7 +29,7 @@ import { getDocsPath } from "../config.js";
import { theme } from "../modes/interactive/theme/theme.js"; import { theme } from "../modes/interactive/theme/theme.js";
import { stripFrontmatter } from "../utils/frontmatter.js"; import { stripFrontmatter } from "../utils/frontmatter.js";
import { sleep } from "../utils/sleep.js"; import { sleep } from "../utils/sleep.js";
import { type BashResult, executeBash as executeBashCommand, executeBashWithOperations } from "./bash-executor.js"; import { type BashResult, executeBashWithOperations } from "./bash-executor.js";
import { import {
type CompactionResult, type CompactionResult,
calculateContextTokens, calculateContextTokens,
@@ -54,9 +54,8 @@ import {
type MessageStartEvent, type MessageStartEvent,
type MessageUpdateEvent, type MessageUpdateEvent,
type SessionBeforeCompactResult, type SessionBeforeCompactResult,
type SessionBeforeForkResult,
type SessionBeforeSwitchResult,
type SessionBeforeTreeResult, type SessionBeforeTreeResult,
type SessionStartEvent,
type ShutdownHandler, type ShutdownHandler,
type ToolDefinition, type ToolDefinition,
type ToolExecutionEndEvent, type ToolExecutionEndEvent,
@@ -78,7 +77,7 @@ import type { SettingsManager } from "./settings-manager.js";
import type { SlashCommandInfo } from "./slash-commands.js"; import type { SlashCommandInfo } from "./slash-commands.js";
import { createSyntheticSourceInfo, type SourceInfo } from "./source-info.js"; import { createSyntheticSourceInfo, type SourceInfo } from "./source-info.js";
import { buildSystemPrompt } from "./system-prompt.js"; import { buildSystemPrompt } from "./system-prompt.js";
import type { BashOperations } from "./tools/bash.js"; import { type BashOperations, createLocalBashOperations } from "./tools/bash.js";
import { createAllToolDefinitions } from "./tools/index.js"; import { createAllToolDefinitions } from "./tools/index.js";
import { createToolDefinitionFromAgentTool, wrapToolDefinition } from "./tools/tool-definition-wrapper.js"; import { createToolDefinitionFromAgentTool, wrapToolDefinition } from "./tools/tool-definition-wrapper.js";
@@ -160,6 +159,8 @@ export interface AgentSessionConfig {
baseToolsOverride?: Record<string, AgentTool>; baseToolsOverride?: Record<string, AgentTool>;
/** Mutable ref used by Agent to access the current ExtensionRunner */ /** Mutable ref used by Agent to access the current ExtensionRunner */
extensionRunnerRef?: { current?: ExtensionRunner }; extensionRunnerRef?: { current?: ExtensionRunner };
/** Session start event metadata emitted when extensions bind to this runtime. */
sessionStartEvent?: SessionStartEvent;
} }
export interface ExtensionBindings { export interface ExtensionBindings {
@@ -276,6 +277,7 @@ export class AgentSession {
private _extensionRunnerRef?: { current?: ExtensionRunner }; private _extensionRunnerRef?: { current?: ExtensionRunner };
private _initialActiveToolNames?: string[]; private _initialActiveToolNames?: string[];
private _baseToolsOverride?: Record<string, AgentTool>; private _baseToolsOverride?: Record<string, AgentTool>;
private _sessionStartEvent: SessionStartEvent;
private _extensionUIContext?: ExtensionUIContext; private _extensionUIContext?: ExtensionUIContext;
private _extensionCommandContextActions?: ExtensionCommandContextActions; private _extensionCommandContextActions?: ExtensionCommandContextActions;
private _extensionShutdownHandler?: ShutdownHandler; private _extensionShutdownHandler?: ShutdownHandler;
@@ -306,6 +308,7 @@ export class AgentSession {
this._extensionRunnerRef = config.extensionRunnerRef; this._extensionRunnerRef = config.extensionRunnerRef;
this._initialActiveToolNames = config.initialActiveToolNames; this._initialActiveToolNames = config.initialActiveToolNames;
this._baseToolsOverride = config.baseToolsOverride; this._baseToolsOverride = config.baseToolsOverride;
this._sessionStartEvent = config.sessionStartEvent ?? { type: "session_start", reason: "startup" };
// Always subscribe to agent events for internal handling // Always subscribe to agent events for internal handling
// (session persistence, extensions, auto-compaction, retry logic) // (session persistence, extensions, auto-compaction, retry logic)
@@ -1346,66 +1349,6 @@ export class AgentSession {
await this.agent.waitForIdle(); await this.agent.waitForIdle();
} }
/**
* Start a new session, optionally with initial messages and parent tracking.
* Clears all messages and starts a new session.
* Listeners are preserved and will continue receiving events.
* @param options.parentSession - Optional parent session path for tracking
* @param options.setup - Optional callback to initialize session (e.g., append messages)
* @returns true if completed, false if cancelled by extension
*/
async newSession(options?: {
parentSession?: string;
setup?: (sessionManager: SessionManager) => Promise<void>;
}): Promise<boolean> {
const previousSessionFile = this.sessionFile;
// Emit session_before_switch event with reason "new" (can be cancelled)
if (this._extensionRunner?.hasHandlers("session_before_switch")) {
const result = (await this._extensionRunner.emit({
type: "session_before_switch",
reason: "new",
})) as SessionBeforeSwitchResult | undefined;
if (result?.cancel) {
return false;
}
}
this._disconnectFromAgent();
await this.abort();
this.agent.reset();
this.sessionManager.newSession({ parentSession: options?.parentSession });
this.agent.sessionId = this.sessionManager.getSessionId();
this._steeringMessages = [];
this._followUpMessages = [];
this._pendingNextTurnMessages = [];
this.sessionManager.appendThinkingLevelChange(this.thinkingLevel);
// Run setup callback if provided (e.g., to append initial messages)
if (options?.setup) {
await options.setup(this.sessionManager);
// Sync agent state with session manager after setup
const sessionContext = this.sessionManager.buildSessionContext();
this.agent.state.messages = sessionContext.messages;
}
this._reconnectToAgent();
// Emit session_switch event with reason "new" to extensions
if (this._extensionRunner) {
await this._extensionRunner.emit({
type: "session_switch",
reason: "new",
previousSessionFile,
});
}
// Emit session event to custom tools
return true;
}
// ========================================================================= // =========================================================================
// Model Management // Model Management
// ========================================================================= // =========================================================================
@@ -1460,12 +1403,8 @@ export class AgentSession {
return this._cycleAvailableModel(direction); return this._cycleAvailableModel(direction);
} }
private _getScopedModelsWithAuth(): Array<{ model: Model<any>; thinkingLevel?: ThinkingLevel }> {
return this._scopedModels.filter((scoped) => this._modelRegistry.hasConfiguredAuth(scoped.model));
}
private async _cycleScopedModel(direction: "forward" | "backward"): Promise<ModelCycleResult | undefined> { private async _cycleScopedModel(direction: "forward" | "backward"): Promise<ModelCycleResult | undefined> {
const scopedModels = this._getScopedModelsWithAuth(); const scopedModels = this._scopedModels.filter((scoped) => this._modelRegistry.hasConfiguredAuth(scoped.model));
if (scopedModels.length <= 1) return undefined; if (scopedModels.length <= 1) return undefined;
const currentModel = this.model; const currentModel = this.model;
@@ -2081,8 +2020,8 @@ export class AgentSession {
if (this._extensionRunner) { if (this._extensionRunner) {
this._applyExtensionBindings(this._extensionRunner); this._applyExtensionBindings(this._extensionRunner);
await this._extensionRunner.emit({ type: "session_start" }); await this._extensionRunner.emit(this._sessionStartEvent);
await this.extendResourcesFromExtensions("startup"); await this.extendResourcesFromExtensions(this._sessionStartEvent.reason === "reload" ? "reload" : "startup");
} }
} }
@@ -2426,7 +2365,7 @@ export class AgentSession {
this._extensionShutdownHandler || this._extensionShutdownHandler ||
this._extensionErrorListener; this._extensionErrorListener;
if (this._extensionRunner && hasBindings) { if (this._extensionRunner && hasBindings) {
await this._extensionRunner.emit({ type: "session_start" }); await this._extensionRunner.emit({ type: "session_start", reason: "reload" });
await this.extendResourcesFromExtensions("reload"); await this.extendResourcesFromExtensions("reload");
} }
} }
@@ -2596,15 +2535,15 @@ export class AgentSession {
const resolvedCommand = prefix ? `${prefix}\n${command}` : command; const resolvedCommand = prefix ? `${prefix}\n${command}` : command;
try { try {
const result = options?.operations const result = await executeBashWithOperations(
? await executeBashWithOperations(resolvedCommand, process.cwd(), options.operations, { resolvedCommand,
onChunk, this.sessionManager.getCwd(),
signal: this._bashAbortController.signal, options?.operations ?? createLocalBashOperations(),
}) {
: await executeBashCommand(resolvedCommand, { onChunk,
onChunk, signal: this._bashAbortController.signal,
signal: this._bashAbortController.signal, },
}); );
this.recordBashResult(command, result, options); this.recordBashResult(command, result, options);
return result; return result;
@@ -2682,86 +2621,6 @@ export class AgentSession {
// Session Management // Session Management
// ========================================================================= // =========================================================================
/**
* Switch to a different session file.
* Aborts current operation, loads messages, restores model/thinking.
* Listeners are preserved and will continue receiving events.
* @returns true if switch completed, false if cancelled by extension
*/
async switchSession(sessionPath: string): Promise<boolean> {
const previousSessionFile = this.sessionManager.getSessionFile();
// Emit session_before_switch event (can be cancelled)
if (this._extensionRunner?.hasHandlers("session_before_switch")) {
const result = (await this._extensionRunner.emit({
type: "session_before_switch",
reason: "resume",
targetSessionFile: sessionPath,
})) as SessionBeforeSwitchResult | undefined;
if (result?.cancel) {
return false;
}
}
this._disconnectFromAgent();
await this.abort();
this._steeringMessages = [];
this._followUpMessages = [];
this._pendingNextTurnMessages = [];
// Set new session
this.sessionManager.setSessionFile(sessionPath);
this.agent.sessionId = this.sessionManager.getSessionId();
// Reload messages
const sessionContext = this.sessionManager.buildSessionContext();
// Emit session_switch event to extensions
if (this._extensionRunner) {
await this._extensionRunner.emit({
type: "session_switch",
reason: "resume",
previousSessionFile,
});
}
// Emit session event to custom tools
this.agent.state.messages = sessionContext.messages;
// Restore model if saved
if (sessionContext.model) {
const previousModel = this.model;
const availableModels = await this._modelRegistry.getAvailable();
const match = availableModels.find(
(m) => m.provider === sessionContext.model!.provider && m.id === sessionContext.model!.modelId,
);
if (match) {
this.agent.state.model = match;
await this._emitModelSelect(match, previousModel, "restore");
}
}
const hasThinkingEntry = this.sessionManager.getBranch().some((entry) => entry.type === "thinking_level_change");
const defaultThinkingLevel = this.settingsManager.getDefaultThinkingLevel() ?? DEFAULT_THINKING_LEVEL;
if (hasThinkingEntry) {
// Restore thinking level if saved (setThinkingLevel clamps to model capabilities)
this.setThinkingLevel(sessionContext.thinkingLevel as ThinkingLevel);
} else {
const availableLevels = this.getAvailableThinkingLevels();
const effectiveLevel = availableLevels.includes(defaultThinkingLevel)
? defaultThinkingLevel
: this._clampThinkingLevel(defaultThinkingLevel, availableLevels);
this.agent.state.thinkingLevel = effectiveLevel;
this.sessionManager.appendThinkingLevelChange(effectiveLevel);
}
this._reconnectToAgent();
return true;
}
/** /**
* Set a display name for the current session. * Set a display name for the current session.
*/ */
@@ -2769,70 +2628,6 @@ export class AgentSession {
this.sessionManager.appendSessionInfo(name); this.sessionManager.appendSessionInfo(name);
} }
/**
* Create a fork from a specific entry.
* Emits before_fork/fork session events to extensions.
*
* @param entryId ID of the entry to fork from
* @returns Object with:
* - selectedText: The text of the selected user message (for editor pre-fill)
* - cancelled: True if an extension cancelled the fork
*/
async fork(entryId: string): Promise<{ selectedText: string; cancelled: boolean }> {
const previousSessionFile = this.sessionFile;
const selectedEntry = this.sessionManager.getEntry(entryId);
if (!selectedEntry || selectedEntry.type !== "message" || selectedEntry.message.role !== "user") {
throw new Error("Invalid entry ID for forking");
}
const selectedText = this._extractUserMessageText(selectedEntry.message.content);
let skipConversationRestore = false;
// Emit session_before_fork event (can be cancelled)
if (this._extensionRunner?.hasHandlers("session_before_fork")) {
const result = (await this._extensionRunner.emit({
type: "session_before_fork",
entryId,
})) as SessionBeforeForkResult | undefined;
if (result?.cancel) {
return { selectedText, cancelled: true };
}
skipConversationRestore = result?.skipConversationRestore ?? false;
}
// Clear pending messages (bound to old session state)
this._pendingNextTurnMessages = [];
if (!selectedEntry.parentId) {
this.sessionManager.newSession({ parentSession: previousSessionFile });
} else {
this.sessionManager.createBranchedSession(selectedEntry.parentId);
}
this.agent.sessionId = this.sessionManager.getSessionId();
// Reload messages from entries (works for both file and in-memory mode)
const sessionContext = this.sessionManager.buildSessionContext();
// Emit session_fork event to extensions (after fork completes)
if (this._extensionRunner) {
await this._extensionRunner.emit({
type: "session_fork",
previousSessionFile,
});
}
// Emit session event to custom tools (with reason "fork")
if (!skipConversationRestore) {
this.agent.state.messages = sessionContext.messages;
}
return { selectedText, cancelled: false };
}
// ========================================================================= // =========================================================================
// Tree Navigation // Tree Navigation
// ========================================================================= // =========================================================================
@@ -3165,6 +2960,7 @@ export class AgentSession {
const toolRenderer: ToolHtmlRenderer = createToolHtmlRenderer({ const toolRenderer: ToolHtmlRenderer = createToolHtmlRenderer({
getToolDefinition: (name) => this.getToolDefinition(name), getToolDefinition: (name) => this.getToolDefinition(name),
theme, theme,
cwd: this.sessionManager.getCwd(),
}); });
return await exportSessionToHtml(this.sessionManager, this.state, { return await exportSessionToHtml(this.sessionManager, this.state, {
@@ -3210,32 +3006,6 @@ export class AgentSession {
return filePath; return filePath;
} }
/**
* Import a JSONL session file.
* Copies the file into the session directory and switches to it (like /resume).
* @param inputPath Path to the JSONL file to import.
* @returns true if the session was switched successfully.
*/
async importFromJsonl(inputPath: string): Promise<boolean> {
const resolved = resolve(inputPath);
if (!existsSync(resolved)) {
throw new Error(`File not found: ${resolved}`);
}
// Copy into the session directory so we don't modify the original
const sessionDir = this.sessionManager.getSessionDir();
if (!existsSync(sessionDir)) {
mkdirSync(sessionDir, { recursive: true });
}
const destPath = join(sessionDir, basename(resolved));
// Avoid overwriting if source and destination are the same file
if (resolve(destPath) !== resolved) {
copyFileSync(resolved, destPath);
}
return this.switchSession(destPath);
}
// ========================================================================= // =========================================================================
// Utilities // Utilities
// ========================================================================= // =========================================================================

View File

@@ -16,6 +16,8 @@ export interface ToolHtmlRendererDeps {
getToolDefinition: (name: string) => ToolDefinition | undefined; getToolDefinition: (name: string) => ToolDefinition | undefined;
/** Theme for styling */ /** Theme for styling */
theme: Theme; theme: Theme;
/** Working directory for render context */
cwd: string;
/** Terminal width for rendering (default: 100) */ /** Terminal width for rendering (default: 100) */
width?: number; width?: number;
} }
@@ -40,7 +42,7 @@ export interface ToolHtmlRenderer {
* methods, converting the resulting TUI Component output (ANSI) to HTML. * methods, converting the resulting TUI Component output (ANSI) to HTML.
*/ */
export function createToolHtmlRenderer(deps: ToolHtmlRendererDeps): ToolHtmlRenderer { export function createToolHtmlRenderer(deps: ToolHtmlRendererDeps): ToolHtmlRenderer {
const { getToolDefinition, theme, width = 100 } = deps; const { getToolDefinition, theme, cwd, width = 100 } = deps;
const renderedCallComponents = new Map<string, Component>(); const renderedCallComponents = new Map<string, Component>();
const renderedResultComponents = new Map<string, Component>(); const renderedResultComponents = new Map<string, Component>();
@@ -69,7 +71,7 @@ export function createToolHtmlRenderer(deps: ToolHtmlRendererDeps): ToolHtmlRend
invalidate: () => {}, invalidate: () => {},
lastComponent, lastComponent,
state: getState(toolCallId), state: getState(toolCallId),
cwd: process.cwd(), cwd,
executionStarted: true, executionStarted: true,
argsComplete: true, argsComplete: true,
isPartial, isPartial,

View File

@@ -120,11 +120,9 @@ export type {
SessionDirectoryHandler, SessionDirectoryHandler,
SessionDirectoryResult, SessionDirectoryResult,
SessionEvent, SessionEvent,
SessionForkEvent,
SessionShutdownEvent, SessionShutdownEvent,
// Events - Session // Events - Session
SessionStartEvent, SessionStartEvent,
SessionSwitchEvent,
SessionTreeEvent, SessionTreeEvent,
SetActiveToolsHandler, SetActiveToolsHandler,
SetLabelHandler, SetLabelHandler,

View File

@@ -432,9 +432,13 @@ export interface SessionDirectoryEvent {
cwd: string; cwd: string;
} }
/** Fired on initial session load */ /** Fired when a session is started, loaded, or reloaded */
export interface SessionStartEvent { export interface SessionStartEvent {
type: "session_start"; type: "session_start";
/** Why this session start happened. */
reason: "startup" | "reload" | "new" | "resume" | "fork";
/** Previously active session file. Present for "new", "resume", and "fork". */
previousSessionFile?: string;
} }
/** Fired before switching to another session (can be cancelled) */ /** Fired before switching to another session (can be cancelled) */
@@ -444,25 +448,12 @@ export interface SessionBeforeSwitchEvent {
targetSessionFile?: string; targetSessionFile?: string;
} }
/** Fired after switching to another session */
export interface SessionSwitchEvent {
type: "session_switch";
reason: "new" | "resume";
previousSessionFile: string | undefined;
}
/** Fired before forking a session (can be cancelled) */ /** Fired before forking a session (can be cancelled) */
export interface SessionBeforeForkEvent { export interface SessionBeforeForkEvent {
type: "session_before_fork"; type: "session_before_fork";
entryId: string; entryId: string;
} }
/** Fired after forking a session */
export interface SessionForkEvent {
type: "session_fork";
previousSessionFile: string | undefined;
}
/** Fired before context compaction (can be cancelled or customized) */ /** Fired before context compaction (can be cancelled or customized) */
export interface SessionBeforeCompactEvent { export interface SessionBeforeCompactEvent {
type: "session_before_compact"; type: "session_before_compact";
@@ -519,9 +510,7 @@ export type SessionEvent =
| SessionDirectoryEvent | SessionDirectoryEvent
| SessionStartEvent | SessionStartEvent
| SessionBeforeSwitchEvent | SessionBeforeSwitchEvent
| SessionSwitchEvent
| SessionBeforeForkEvent | SessionBeforeForkEvent
| SessionForkEvent
| SessionBeforeCompactEvent | SessionBeforeCompactEvent
| SessionCompactEvent | SessionCompactEvent
| SessionShutdownEvent | SessionShutdownEvent
@@ -1008,9 +997,7 @@ export interface ExtensionAPI {
event: "session_before_switch", event: "session_before_switch",
handler: ExtensionHandler<SessionBeforeSwitchEvent, SessionBeforeSwitchResult>, handler: ExtensionHandler<SessionBeforeSwitchEvent, SessionBeforeSwitchResult>,
): void; ): void;
on(event: "session_switch", handler: ExtensionHandler<SessionSwitchEvent>): void;
on(event: "session_before_fork", handler: ExtensionHandler<SessionBeforeForkEvent, SessionBeforeForkResult>): void; on(event: "session_before_fork", handler: ExtensionHandler<SessionBeforeForkEvent, SessionBeforeForkResult>): void;
on(event: "session_fork", handler: ExtensionHandler<SessionForkEvent>): void;
on( on(
event: "session_before_compact", event: "session_before_compact",
handler: ExtensionHandler<SessionBeforeCompactEvent, SessionBeforeCompactResult>, handler: ExtensionHandler<SessionBeforeCompactEvent, SessionBeforeCompactResult>,

View File

@@ -1,5 +1,5 @@
import { type ExecFileException, execFile, spawnSync } from "child_process"; import { type ExecFileException, execFile, spawnSync } from "child_process";
import { existsSync, type FSWatcher, readFileSync, statSync, watch } from "fs"; import { existsSync, type FSWatcher, readFileSync, statSync, unwatchFile, watch, watchFile } from "fs";
import { dirname, join, resolve } from "path"; import { dirname, join, resolve } from "path";
type GitPaths = { type GitPaths = {
@@ -12,8 +12,8 @@ type GitPaths = {
* Find git metadata paths by walking up from cwd. * Find git metadata paths by walking up from cwd.
* Handles both regular git repos (.git is a directory) and worktrees (.git is a file). * Handles both regular git repos (.git is a directory) and worktrees (.git is a file).
*/ */
function findGitPaths(): GitPaths | null { function findGitPaths(cwd: string): GitPaths | null {
let dir = process.cwd(); let dir = cwd;
while (true) { while (true) {
const gitPath = join(dir, ".git"); const gitPath = join(dir, ".git");
if (existsSync(gitPath)) { if (existsSync(gitPath)) {
@@ -84,6 +84,7 @@ function resolveBranchWithGitAsync(repoDir: string): Promise<string | null> {
* Token stats, model info available via ctx.sessionManager and ctx.model. * Token stats, model info available via ctx.sessionManager and ctx.model.
*/ */
export class FooterDataProvider { export class FooterDataProvider {
private cwd: string;
private static readonly WATCH_DEBOUNCE_MS = 500; private static readonly WATCH_DEBOUNCE_MS = 500;
private extensionStatuses = new Map<string, string>(); private extensionStatuses = new Map<string, string>();
@@ -91,6 +92,8 @@ export class FooterDataProvider {
private gitPaths: GitPaths | null | undefined = undefined; private gitPaths: GitPaths | null | undefined = undefined;
private headWatcher: FSWatcher | null = null; private headWatcher: FSWatcher | null = null;
private reftableWatcher: FSWatcher | null = null; private reftableWatcher: FSWatcher | null = null;
private reftableTablesListWatcher: FSWatcher | null = null;
private reftableTablesListPath: string | null = null;
private branchChangeCallbacks = new Set<() => void>(); private branchChangeCallbacks = new Set<() => void>();
private availableProviderCount = 0; private availableProviderCount = 0;
private refreshTimer: ReturnType<typeof setTimeout> | null = null; private refreshTimer: ReturnType<typeof setTimeout> | null = null;
@@ -98,8 +101,9 @@ export class FooterDataProvider {
private refreshPending = false; private refreshPending = false;
private disposed = false; private disposed = false;
constructor() { constructor(cwd: string = process.cwd()) {
this.gitPaths = findGitPaths(); this.cwd = cwd;
this.gitPaths = findGitPaths(cwd);
this.setupGitWatcher(); this.setupGitWatcher();
} }
@@ -146,6 +150,38 @@ export class FooterDataProvider {
this.availableProviderCount = count; this.availableProviderCount = count;
} }
setCwd(cwd: string): void {
if (this.cwd === cwd) {
return;
}
this.cwd = cwd;
if (this.refreshTimer) {
clearTimeout(this.refreshTimer);
this.refreshTimer = null;
}
if (this.headWatcher) {
this.headWatcher.close();
this.headWatcher = null;
}
if (this.reftableWatcher) {
this.reftableWatcher.close();
this.reftableWatcher = null;
}
if (this.reftableTablesListWatcher) {
this.reftableTablesListWatcher.close();
this.reftableTablesListWatcher = null;
}
if (this.reftableTablesListPath) {
unwatchFile(this.reftableTablesListPath);
this.reftableTablesListPath = null;
}
this.cachedBranch = undefined;
this.gitPaths = findGitPaths(cwd);
this.setupGitWatcher();
this.notifyBranchChange();
}
/** Internal: cleanup */ /** Internal: cleanup */
dispose(): void { dispose(): void {
this.disposed = true; this.disposed = true;
@@ -161,6 +197,14 @@ export class FooterDataProvider {
this.reftableWatcher.close(); this.reftableWatcher.close();
this.reftableWatcher = null; this.reftableWatcher = null;
} }
if (this.reftableTablesListWatcher) {
this.reftableTablesListWatcher.close();
this.reftableTablesListWatcher = null;
}
if (this.reftableTablesListPath) {
unwatchFile(this.reftableTablesListPath);
this.reftableTablesListPath = null;
}
this.branchChangeCallbacks.clear(); this.branchChangeCallbacks.clear();
} }
@@ -169,9 +213,10 @@ export class FooterDataProvider {
} }
private scheduleRefresh(): void { private scheduleRefresh(): void {
if (this.disposed) return; if (this.disposed || this.refreshTimer) return;
if (this.refreshTimer) { if (this.refreshInFlight) {
clearTimeout(this.refreshTimer); this.refreshPending = true;
return;
} }
this.refreshTimer = setTimeout(() => { this.refreshTimer = setTimeout(() => {
this.refreshTimer = null; this.refreshTimer = null;
@@ -262,6 +307,27 @@ export class FooterDataProvider {
} catch { } catch {
// Silently fail if we can't watch // Silently fail if we can't watch
} }
const tablesListPath = join(reftableDir, "tables.list");
if (existsSync(tablesListPath)) {
this.reftableTablesListPath = tablesListPath;
try {
this.reftableTablesListWatcher = watch(tablesListPath, () => {
this.scheduleRefresh();
});
} catch {
// Silently fail if we can't watch
}
watchFile(tablesListPath, { interval: 250 }, (current, previous) => {
if (
current.mtimeMs !== previous.mtimeMs ||
current.ctimeMs !== previous.ctimeMs ||
current.size !== previous.size
) {
this.scheduleRefresh();
}
});
}
} }
} }
} }

View File

@@ -11,6 +11,12 @@ export {
type PromptOptions, type PromptOptions,
type SessionStats, type SessionStats,
} from "./agent-session.js"; } from "./agent-session.js";
export {
type AgentSessionRuntimeBootstrap,
AgentSessionRuntimeHost,
type CreateAgentSessionRuntimeOptions,
createAgentSessionRuntime,
} from "./agent-session-runtime.js";
export { type BashExecutorOptions, type BashResult, executeBash, executeBashWithOperations } from "./bash-executor.js"; export { type BashExecutorOptions, type BashResult, executeBash, executeBashWithOperations } from "./bash-executor.js";
export type { CompactionResult } from "./compaction/index.js"; export type { CompactionResult } from "./compaction/index.js";
export { createEventBus, type EventBus, type EventBusController } from "./event-bus.js"; export { createEventBus, type EventBus, type EventBusController } from "./event-bus.js";
@@ -45,10 +51,8 @@ export {
type SessionBeforeSwitchEvent, type SessionBeforeSwitchEvent,
type SessionBeforeTreeEvent, type SessionBeforeTreeEvent,
type SessionCompactEvent, type SessionCompactEvent,
type SessionForkEvent,
type SessionShutdownEvent, type SessionShutdownEvent,
type SessionStartEvent, type SessionStartEvent,
type SessionSwitchEvent,
type SessionTreeEvent, type SessionTreeEvent,
type ToolCallEvent, type ToolCallEvent,
type ToolCallEventResult, type ToolCallEventResult,

View File

@@ -5,7 +5,7 @@ import { getAgentDir, getDocsPath } from "../config.js";
import { AgentSession } from "./agent-session.js"; import { AgentSession } from "./agent-session.js";
import { AuthStorage } from "./auth-storage.js"; import { AuthStorage } from "./auth-storage.js";
import { DEFAULT_THINKING_LEVEL } from "./defaults.js"; import { DEFAULT_THINKING_LEVEL } from "./defaults.js";
import type { ExtensionRunner, LoadExtensionsResult, ToolDefinition } from "./extensions/index.js"; import type { ExtensionRunner, LoadExtensionsResult, SessionStartEvent, ToolDefinition } from "./extensions/index.js";
import { convertToLlm } from "./messages.js"; import { convertToLlm } from "./messages.js";
import { ModelRegistry } from "./model-registry.js"; import { ModelRegistry } from "./model-registry.js";
import { findInitialModel } from "./model-resolver.js"; import { findInitialModel } from "./model-resolver.js";
@@ -70,6 +70,8 @@ export interface CreateAgentSessionOptions {
/** Settings manager. Default: SettingsManager.create(cwd, agentDir) */ /** Settings manager. Default: SettingsManager.create(cwd, agentDir) */
settingsManager?: SettingsManager; settingsManager?: SettingsManager;
/** Session start event metadata for extension runtime startup. */
sessionStartEvent?: SessionStartEvent;
} }
/** Result from createAgentSession */ /** Result from createAgentSession */
@@ -84,6 +86,12 @@ export interface CreateAgentSessionResult {
// Re-exports // Re-exports
export {
type AgentSessionRuntimeBootstrap,
AgentSessionRuntimeHost,
type CreateAgentSessionRuntimeOptions,
createAgentSessionRuntime,
} from "./agent-session-runtime.js";
export type { export type {
ExtensionAPI, ExtensionAPI,
ExtensionCommandContext, ExtensionCommandContext,
@@ -349,6 +357,7 @@ export async function createAgentSession(options: CreateAgentSessionOptions = {}
modelRegistry, modelRegistry,
initialActiveToolNames, initialActiveToolNames,
extensionRunnerRef, extensionRunnerRef,
sessionStartEvent: options.sessionStartEvent,
}); });
const extensionsResult = resourceLoader.getExtensions(); const extensionsResult = resourceLoader.getExtensions();

View File

@@ -103,10 +103,8 @@ export type {
SessionBeforeSwitchEvent, SessionBeforeSwitchEvent,
SessionBeforeTreeEvent, SessionBeforeTreeEvent,
SessionCompactEvent, SessionCompactEvent,
SessionForkEvent,
SessionShutdownEvent, SessionShutdownEvent,
SessionStartEvent, SessionStartEvent,
SessionSwitchEvent,
SessionTreeEvent, SessionTreeEvent,
SlashCommandInfo, SlashCommandInfo,
SlashCommandSource, SlashCommandSource,
@@ -157,10 +155,14 @@ export type { ResourceCollision, ResourceDiagnostic, ResourceLoader } from "./co
export { DefaultResourceLoader } from "./core/resource-loader.js"; export { DefaultResourceLoader } from "./core/resource-loader.js";
// SDK for programmatic usage // SDK for programmatic usage
export { export {
type AgentSessionRuntimeBootstrap,
AgentSessionRuntimeHost,
type CreateAgentSessionOptions, type CreateAgentSessionOptions,
type CreateAgentSessionResult, type CreateAgentSessionResult,
type CreateAgentSessionRuntimeOptions,
// Factory // Factory
createAgentSession, createAgentSession,
createAgentSessionRuntime,
createBashTool, createBashTool,
// Tool factories (for custom cwd) // Tool factories (for custom cwd)
createCodingTools, createCodingTools,

View File

@@ -5,6 +5,7 @@
* createAgentSession() options. The SDK does the heavy lifting. * createAgentSession() options. The SDK does the heavy lifting.
*/ */
import { resolve } from "node:path";
import { type ImageContent, modelsAreEqual, supportsXhigh } from "@mariozechner/pi-ai"; import { type ImageContent, modelsAreEqual, supportsXhigh } from "@mariozechner/pi-ai";
import chalk from "chalk"; import chalk from "chalk";
import { createInterface } from "readline"; import { createInterface } from "readline";
@@ -15,6 +16,11 @@ import { buildInitialMessage } from "./cli/initial-message.js";
import { listModels } from "./cli/list-models.js"; import { listModels } from "./cli/list-models.js";
import { selectSession } from "./cli/session-picker.js"; import { selectSession } from "./cli/session-picker.js";
import { APP_NAME, getAgentDir, getModelsPath, VERSION } from "./config.js"; import { APP_NAME, getAgentDir, getModelsPath, VERSION } from "./config.js";
import {
type AgentSessionRuntimeBootstrap,
AgentSessionRuntimeHost,
createAgentSessionRuntime,
} from "./core/agent-session-runtime.js";
import { AuthStorage } from "./core/auth-storage.js"; import { AuthStorage } from "./core/auth-storage.js";
import { exportFromFile } from "./core/export-html/index.js"; import { exportFromFile } from "./core/export-html/index.js";
import type { LoadExtensionsResult } from "./core/extensions/index.js"; import type { LoadExtensionsResult } from "./core/extensions/index.js";
@@ -24,7 +30,7 @@ import { resolveCliModel, resolveModelScope, type ScopedModel } from "./core/mod
import { restoreStdout, takeOverStdout } from "./core/output-guard.js"; import { restoreStdout, takeOverStdout } from "./core/output-guard.js";
import { DefaultPackageManager } from "./core/package-manager.js"; import { DefaultPackageManager } from "./core/package-manager.js";
import { DefaultResourceLoader } from "./core/resource-loader.js"; import { DefaultResourceLoader } from "./core/resource-loader.js";
import { type CreateAgentSessionOptions, createAgentSession } from "./core/sdk.js"; import type { CreateAgentSessionOptions } from "./core/sdk.js";
import { SessionManager } from "./core/session-manager.js"; import { SessionManager } from "./core/session-manager.js";
import { SettingsManager } from "./core/settings-manager.js"; import { SettingsManager } from "./core/settings-manager.js";
import { printTimings, resetTimings, time } from "./core/timings.js"; import { printTimings, resetTimings, time } from "./core/timings.js";
@@ -597,6 +603,40 @@ function buildSessionOptions(
return { options, cliThinkingFromModel }; return { options, cliThinkingFromModel };
} }
function resolveCliPaths(cwd: string, paths: string[] | undefined): string[] | undefined {
return paths?.map((value) => resolve(cwd, value));
}
function buildRuntimeBootstrap(
parsed: Args,
cwd: string,
agentDir: string,
authStorage: AuthStorage,
sessionOptions: CreateAgentSessionOptions,
): AgentSessionRuntimeBootstrap {
return {
agentDir,
authStorage,
model: sessionOptions.model,
thinkingLevel: sessionOptions.thinkingLevel,
scopedModels: sessionOptions.scopedModels,
tools: sessionOptions.tools,
customTools: sessionOptions.customTools,
resourceLoader: {
additionalExtensionPaths: resolveCliPaths(cwd, parsed.extensions),
additionalSkillPaths: resolveCliPaths(cwd, parsed.skills),
additionalPromptTemplatePaths: resolveCliPaths(cwd, parsed.promptTemplates),
additionalThemePaths: resolveCliPaths(cwd, parsed.themes),
noExtensions: parsed.noExtensions,
noSkills: parsed.noSkills,
noPromptTemplates: parsed.noPromptTemplates,
noThemes: parsed.noThemes,
systemPrompt: parsed.systemPrompt,
appendSystemPrompt: parsed.appendSystemPrompt,
},
};
}
async function handleConfigCommand(args: string[]): Promise<boolean> { async function handleConfigCommand(args: string[]): Promise<boolean> {
if (args[0] !== "config") { if (args[0] !== "config") {
return false; return false;
@@ -818,11 +858,7 @@ export async function main(args: string[]) {
modelRegistry, modelRegistry,
settingsManager, settingsManager,
); );
sessionOptions.authStorage = authStorage;
sessionOptions.modelRegistry = modelRegistry;
sessionOptions.resourceLoader = resourceLoader;
// Handle CLI --api-key as runtime override (not persisted)
if (parsed.apiKey) { if (parsed.apiKey) {
if (!sessionOptions.model) { if (!sessionOptions.model) {
console.error( console.error(
@@ -833,7 +869,16 @@ export async function main(args: string[]) {
authStorage.setRuntimeApiKey(sessionOptions.model.provider, parsed.apiKey); authStorage.setRuntimeApiKey(sessionOptions.model.provider, parsed.apiKey);
} }
const { session, modelFallbackMessage } = await createAgentSession(sessionOptions); const runtimeBootstrap = buildRuntimeBootstrap(parsed, cwd, agentDir, authStorage, sessionOptions);
const runtime = await createAgentSessionRuntime(runtimeBootstrap, {
cwd: sessionManager?.getCwd() ?? cwd,
sessionManager,
});
if (process.cwd() !== runtime.cwd) {
process.chdir(runtime.cwd);
}
const runtimeHost = new AgentSessionRuntimeHost(runtimeBootstrap, runtime);
const { session, modelFallbackMessage } = runtime;
time("createAgentSession"); time("createAgentSession");
if (!isInteractive && !session.model) { if (!isInteractive && !session.model) {
@@ -861,7 +906,7 @@ export async function main(args: string[]) {
if (mode === "rpc") { if (mode === "rpc") {
printTimings(); printTimings();
await runRpcMode(session); await runRpcMode(runtimeHost);
} else if (isInteractive) { } else if (isInteractive) {
if (scopedModels.length > 0 && (parsed.verbose || !settingsManager.getQuietStartup())) { if (scopedModels.length > 0 && (parsed.verbose || !settingsManager.getQuietStartup())) {
const modelList = scopedModels const modelList = scopedModels
@@ -873,7 +918,7 @@ export async function main(args: string[]) {
console.log(chalk.dim(`Model scope: ${modelList} ${chalk.gray("(Ctrl+P to cycle)")}`)); console.log(chalk.dim(`Model scope: ${modelList} ${chalk.gray("(Ctrl+P to cycle)")}`));
} }
const interactiveMode = new InteractiveMode(session, { const interactiveMode = new InteractiveMode(runtimeHost, {
migratedProviders, migratedProviders,
modelFallbackMessage, modelFallbackMessage,
initialMessage, initialMessage,
@@ -900,7 +945,7 @@ export async function main(args: string[]) {
await interactiveMode.run(); await interactiveMode.run();
} else { } else {
printTimings(); printTimings();
const exitCode = await runPrintMode(session, { const exitCode = await runPrintMode(runtimeHost, {
mode, mode,
messages: parsed.messages, messages: parsed.messages,
initialMessage, initialMessage,

View File

@@ -38,6 +38,10 @@ export class FooterComponent implements Component {
private footerData: ReadonlyFooterDataProvider, private footerData: ReadonlyFooterDataProvider,
) {} ) {}
setSession(session: AgentSession): void {
this.session = session;
}
setAutoCompactEnabled(enabled: boolean): void { setAutoCompactEnabled(enabled: boolean): void {
this.autoCompactEnabled = enabled; this.autoCompactEnabled = enabled;
} }
@@ -86,7 +90,7 @@ export class FooterComponent implements Component {
const contextPercent = contextUsage?.percent !== null ? contextPercentValue.toFixed(1) : "?"; const contextPercent = contextUsage?.percent !== null ? contextPercentValue.toFixed(1) : "?";
// Replace home directory with ~ // Replace home directory with ~
let pwd = process.cwd(); let pwd = this.session.sessionManager.getCwd();
const home = process.env.HOME || process.env.USERPROFILE; const home = process.env.HOME || process.env.USERPROFILE;
if (home && pwd.startsWith(home)) { if (home && pwd.startsWith(home)) {
pwd = `~${pwd.slice(home.length)}`; pwd = `~${pwd.slice(home.length)}`;

View File

@@ -47,6 +47,7 @@ import {
VERSION, VERSION,
} from "../../config.js"; } from "../../config.js";
import { type AgentSession, type AgentSessionEvent, parseSkillBlock } from "../../core/agent-session.js"; import { type AgentSession, type AgentSessionEvent, parseSkillBlock } from "../../core/agent-session.js";
import type { AgentSessionRuntimeHost } from "../../core/agent-session-runtime.js";
import type { import type {
ExtensionContext, ExtensionContext,
ExtensionRunner, ExtensionRunner,
@@ -144,7 +145,7 @@ export interface InteractiveModeOptions {
} }
export class InteractiveMode { export class InteractiveMode {
private session: AgentSession; private runtimeHost: AgentSessionRuntimeHost;
private ui: TUI; private ui: TUI;
private chatContainer: Container; private chatContainer: Container;
private pendingMessagesContainer: Container; private pendingMessagesContainer: Container;
@@ -242,6 +243,9 @@ export class InteractiveMode {
private customHeader: (Component & { dispose?(): void }) | undefined = undefined; private customHeader: (Component & { dispose?(): void }) | undefined = undefined;
// Convenience accessors // Convenience accessors
private get session(): AgentSession {
return this.runtimeHost.session;
}
private get agent() { private get agent() {
return this.session.agent; return this.session.agent;
} }
@@ -253,10 +257,10 @@ export class InteractiveMode {
} }
constructor( constructor(
session: AgentSession, runtimeHost: AgentSessionRuntimeHost,
private options: InteractiveModeOptions = {}, private options: InteractiveModeOptions = {},
) { ) {
this.session = session; this.runtimeHost = runtimeHost;
this.version = VERSION; this.version = VERSION;
this.ui = new TUI(new ProcessTerminal(), this.settingsManager.getShowHardwareCursor()); this.ui = new TUI(new ProcessTerminal(), this.settingsManager.getShowHardwareCursor());
this.ui.setClearOnShrink(this.settingsManager.getClearOnShrink()); this.ui.setClearOnShrink(this.settingsManager.getClearOnShrink());
@@ -277,9 +281,9 @@ export class InteractiveMode {
this.editor = this.defaultEditor; this.editor = this.defaultEditor;
this.editorContainer = new Container(); this.editorContainer = new Container();
this.editorContainer.addChild(this.editor as Component); this.editorContainer.addChild(this.editor as Component);
this.footerDataProvider = new FooterDataProvider(); this.footerDataProvider = new FooterDataProvider(this.sessionManager.getCwd());
this.footer = new FooterComponent(session, this.footerDataProvider); this.footer = new FooterComponent(this.session, this.footerDataProvider);
this.footer.setAutoCompactEnabled(session.autoCompactionEnabled); this.footer.setAutoCompactEnabled(this.session.autoCompactionEnabled);
// Load hide thinking block setting // Load hide thinking block setting
this.hideThinkingBlock = this.settingsManager.getHideThinkingBlock(); this.hideThinkingBlock = this.settingsManager.getHideThinkingBlock();
@@ -412,7 +416,7 @@ export class InteractiveMode {
// Setup autocomplete // Setup autocomplete
this.autocompleteProvider = new CombinedAutocompleteProvider( this.autocompleteProvider = new CombinedAutocompleteProvider(
[...slashCommands, ...templateCommands, ...extensionCommands, ...skillCommandList], [...slashCommands, ...templateCommands, ...extensionCommands, ...skillCommandList],
process.cwd(), this.sessionManager.getCwd(),
fdPath, fdPath,
); );
this.defaultEditor.setAutocompleteProvider(this.autocompleteProvider); this.defaultEditor.setAutocompleteProvider(this.autocompleteProvider);
@@ -524,7 +528,7 @@ export class InteractiveMode {
this.isInitialized = true; this.isInitialized = true;
// Initialize extensions first so resources are shown before messages // Initialize extensions first so resources are shown before messages
await this.initExtensions(); await this.bindCurrentSessionExtensions();
// Render initial messages AFTER showing loaded resources // Render initial messages AFTER showing loaded resources
this.renderInitialMessages(); this.renderInitialMessages();
@@ -555,7 +559,7 @@ export class InteractiveMode {
* Update terminal title with session name and cwd. * Update terminal title with session name and cwd.
*/ */
private updateTerminalTitle(): void { private updateTerminalTitle(): void {
const cwdBasename = path.basename(process.cwd()); const cwdBasename = path.basename(this.sessionManager.getCwd());
const sessionName = this.sessionManager.getSessionName(); const sessionName = this.sessionManager.getSessionName();
if (sessionName) { if (sessionName) {
this.ui.terminal.setTitle(`π - ${sessionName} - ${cwdBasename}`); this.ui.terminal.setTitle(`π - ${sessionName} - ${cwdBasename}`);
@@ -673,7 +677,7 @@ export class InteractiveMode {
try { try {
const packageManager = new DefaultPackageManager({ const packageManager = new DefaultPackageManager({
cwd: process.cwd(), cwd: this.sessionManager.getCwd(),
agentDir: getAgentDir(), agentDir: getAgentDir(),
settingsManager: this.settingsManager, settingsManager: this.settingsManager,
}); });
@@ -1161,7 +1165,7 @@ export class InteractiveMode {
/** /**
* Initialize the extension system with TUI-based UI context. * Initialize the extension system with TUI-based UI context.
*/ */
private async initExtensions(): Promise<void> { private async bindCurrentSessionExtensions(): Promise<void> {
const uiContext = this.createExtensionUIContext(); const uiContext = this.createExtensionUIContext();
await this.session.bindExtensions({ await this.session.bindExtensions({
uiContext, uiContext,
@@ -1173,39 +1177,23 @@ export class InteractiveMode {
this.loadingAnimation = undefined; this.loadingAnimation = undefined;
} }
this.statusContainer.clear(); this.statusContainer.clear();
const result = await this.runtimeHost.newSession(options);
// Delegate to AgentSession (handles setup + agent state sync) if (!result.cancelled) {
const success = await this.session.newSession(options); await this.handleRuntimeSessionChange();
if (!success) { this.renderCurrentSessionState();
return { cancelled: true }; this.ui.requestRender();
} }
return result;
// Clear UI state
this.chatContainer.clear();
this.pendingMessagesContainer.clear();
this.compactionQueuedMessages = [];
this.streamingComponent = undefined;
this.streamingMessage = undefined;
this.pendingTools.clear();
// Render any messages added via setup, or show empty session
this.renderInitialMessages();
this.ui.requestRender();
return { cancelled: false };
}, },
fork: async (entryId) => { fork: async (entryId) => {
const result = await this.session.fork(entryId); const result = await this.runtimeHost.fork(entryId);
if (result.cancelled) { if (!result.cancelled) {
return { cancelled: true }; await this.handleRuntimeSessionChange();
this.renderCurrentSessionState();
this.editor.setText(result.selectedText ?? "");
this.showStatus("Forked to new session");
} }
return { cancelled: result.cancelled };
this.chatContainer.clear();
this.renderInitialMessages();
this.editor.setText(result.selectedText);
this.showStatus("Forked to new session");
return { cancelled: false };
}, },
navigateTree: async (targetId, options) => { navigateTree: async (targetId, options) => {
const result = await this.session.navigateTree(targetId, { const result = await this.session.navigateTree(targetId, {
@@ -1224,7 +1212,6 @@ export class InteractiveMode {
this.editor.setText(result.editorText); this.editor.setText(result.editorText);
} }
this.showStatus("Navigated to selected point"); this.showStatus("Navigated to selected point");
return { cancelled: false }; return { cancelled: false };
}, },
switchSession: async (sessionPath) => { switchSession: async (sessionPath) => {
@@ -1259,6 +1246,45 @@ export class InteractiveMode {
this.showLoadedResources({ force: false }); this.showLoadedResources({ force: false });
} }
private applyRuntimeSettings(): void {
this.footer.setSession(this.session);
this.footer.setAutoCompactEnabled(this.session.autoCompactionEnabled);
this.footerDataProvider.setCwd(this.sessionManager.getCwd());
this.hideThinkingBlock = this.settingsManager.getHideThinkingBlock();
this.ui.setShowHardwareCursor(this.settingsManager.getShowHardwareCursor());
this.ui.setClearOnShrink(this.settingsManager.getClearOnShrink());
const editorPaddingX = this.settingsManager.getEditorPaddingX();
const autocompleteMaxVisible = this.settingsManager.getAutocompleteMaxVisible();
this.defaultEditor.setPaddingX(editorPaddingX);
this.defaultEditor.setAutocompleteMaxVisible(autocompleteMaxVisible);
if (this.editor !== this.defaultEditor) {
this.editor.setPaddingX?.(editorPaddingX);
this.editor.setAutocompleteMaxVisible?.(autocompleteMaxVisible);
}
}
private async handleRuntimeSessionChange(): Promise<void> {
this.resetExtensionUI();
this.unsubscribe?.();
this.unsubscribe = undefined;
this.applyRuntimeSettings();
await this.bindCurrentSessionExtensions();
this.subscribeToAgent();
await this.updateAvailableProviderCount();
this.updateEditorBorderColor();
this.updateTerminalTitle();
}
private renderCurrentSessionState(): void {
this.chatContainer.clear();
this.pendingMessagesContainer.clear();
this.compactionQueuedMessages = [];
this.streamingComponent = undefined;
this.streamingMessage = undefined;
this.pendingTools.clear();
this.renderInitialMessages();
}
/** /**
* Get a registered tool definition by name (for custom rendering). * Get a registered tool definition by name (for custom rendering).
*/ */
@@ -1277,7 +1303,7 @@ export class InteractiveMode {
const createContext = (): ExtensionContext => ({ const createContext = (): ExtensionContext => ({
ui: this.createExtensionUIContext(), ui: this.createExtensionUIContext(),
hasUI: true, hasUI: true,
cwd: process.cwd(), cwd: this.sessionManager.getCwd(),
sessionManager: this.sessionManager, sessionManager: this.sessionManager,
modelRegistry: this.session.modelRegistry, modelRegistry: this.session.modelRegistry,
model: this.session.model, model: this.session.model,
@@ -2305,6 +2331,7 @@ export class InteractiveMode {
}, },
this.getRegisteredToolDefinition(content.name), this.getRegisteredToolDefinition(content.name),
this.ui, this.ui,
this.sessionManager.getCwd(),
); );
component.setExpanded(this.toolOutputExpanded); component.setExpanded(this.toolOutputExpanded);
this.chatContainer.addChild(component); this.chatContainer.addChild(component);
@@ -2372,6 +2399,7 @@ export class InteractiveMode {
}, },
this.getRegisteredToolDefinition(event.toolName), this.getRegisteredToolDefinition(event.toolName),
this.ui, this.ui,
this.sessionManager.getCwd(),
); );
component.setExpanded(this.toolOutputExpanded); component.setExpanded(this.toolOutputExpanded);
this.chatContainer.addChild(component); this.chatContainer.addChild(component);
@@ -2681,6 +2709,7 @@ export class InteractiveMode {
{ showImages: this.settingsManager.getShowImages() }, { showImages: this.settingsManager.getShowImages() },
this.getRegisteredToolDefinition(content.name), this.getRegisteredToolDefinition(content.name),
this.ui, this.ui,
this.sessionManager.getCwd(),
); );
component.setExpanded(this.toolOutputExpanded); component.setExpanded(this.toolOutputExpanded);
this.chatContainer.addChild(component); this.chatContainer.addChild(component);
@@ -2779,14 +2808,7 @@ export class InteractiveMode {
private async shutdown(): Promise<void> { private async shutdown(): Promise<void> {
if (this.isShuttingDown) return; if (this.isShuttingDown) return;
this.isShuttingDown = true; this.isShuttingDown = true;
await this.runtimeHost.dispose();
// Emit shutdown event to extensions
const extensionRunner = this.session.extensionRunner;
if (extensionRunner?.hasHandlers("session_shutdown")) {
await extensionRunner.emit({
type: "session_shutdown",
});
}
// Wait for any pending renders to complete // Wait for any pending renders to complete
// requestRender() uses process.nextTick(), so we wait one tick // requestRender() uses process.nextTick(), so we wait one tick
@@ -3605,17 +3627,15 @@ export class InteractiveMode {
const selector = new UserMessageSelectorComponent( const selector = new UserMessageSelectorComponent(
userMessages.map((m) => ({ id: m.entryId, text: m.text })), userMessages.map((m) => ({ id: m.entryId, text: m.text })),
async (entryId) => { async (entryId) => {
const result = await this.session.fork(entryId); const result = await this.runtimeHost.fork(entryId);
if (result.cancelled) { if (result.cancelled) {
// Extension cancelled the fork
done(); done();
this.ui.requestRender(); this.ui.requestRender();
return; return;
} }
await this.handleRuntimeSessionChange();
this.chatContainer.clear(); this.renderCurrentSessionState();
this.renderInitialMessages(); this.editor.setText(result.selectedText ?? "");
this.editor.setText(result.selectedText);
done(); done();
this.showStatus("Branched to new session"); this.showStatus("Branched to new session");
}, },
@@ -3792,26 +3812,17 @@ export class InteractiveMode {
} }
private async handleResumeSession(sessionPath: string): Promise<void> { private async handleResumeSession(sessionPath: string): Promise<void> {
// Stop loading animation
if (this.loadingAnimation) { if (this.loadingAnimation) {
this.loadingAnimation.stop(); this.loadingAnimation.stop();
this.loadingAnimation = undefined; this.loadingAnimation = undefined;
} }
this.statusContainer.clear(); this.statusContainer.clear();
const result = await this.runtimeHost.switchSession(sessionPath);
// Clear UI state if (result.cancelled) {
this.pendingMessagesContainer.clear(); return;
this.compactionQueuedMessages = []; }
this.streamingComponent = undefined; await this.handleRuntimeSessionChange();
this.streamingMessage = undefined; this.renderCurrentSessionState();
this.pendingTools.clear();
// Switch session via AgentSession (emits extension session events)
await this.session.switchSession(sessionPath);
// Clear and re-render the chat
this.chatContainer.clear();
this.renderInitialMessages();
this.showStatus("Resumed session"); this.showStatus("Resumed session");
} }
@@ -4063,29 +4074,18 @@ export class InteractiveMode {
} }
try { try {
// Stop loading animation
if (this.loadingAnimation) { if (this.loadingAnimation) {
this.loadingAnimation.stop(); this.loadingAnimation.stop();
this.loadingAnimation = undefined; this.loadingAnimation = undefined;
} }
this.statusContainer.clear(); this.statusContainer.clear();
const result = await this.runtimeHost.importFromJsonl(inputPath);
// Clear UI state if (result.cancelled) {
this.pendingMessagesContainer.clear(); this.showStatus("Import cancelled");
this.compactionQueuedMessages = [];
this.streamingComponent = undefined;
this.streamingMessage = undefined;
this.pendingTools.clear();
const success = await this.session.importFromJsonl(inputPath);
if (!success) {
this.showWarning("Import cancelled");
return; return;
} }
await this.handleRuntimeSessionChange();
// Clear and re-render the chat this.renderCurrentSessionState();
this.chatContainer.clear();
this.renderInitialMessages();
this.showStatus(`Session imported from: ${inputPath}`); this.showStatus(`Session imported from: ${inputPath}`);
} catch (error: unknown) { } catch (error: unknown) {
this.showError(`Failed to import session: ${error instanceof Error ? error.message : "Unknown error"}`); this.showError(`Failed to import session: ${error instanceof Error ? error.message : "Unknown error"}`);
@@ -4427,25 +4427,17 @@ export class InteractiveMode {
} }
private async handleClearCommand(): Promise<void> { private async handleClearCommand(): Promise<void> {
// Stop loading animation
if (this.loadingAnimation) { if (this.loadingAnimation) {
this.loadingAnimation.stop(); this.loadingAnimation.stop();
this.loadingAnimation = undefined; this.loadingAnimation = undefined;
} }
this.statusContainer.clear(); this.statusContainer.clear();
const result = await this.runtimeHost.newSession();
// New session via session (emits extension session events) if (result.cancelled) {
await this.session.newSession(); return;
}
// Clear UI state await this.handleRuntimeSessionChange();
this.headerContainer.clear(); this.renderCurrentSessionState();
this.chatContainer.clear();
this.pendingMessagesContainer.clear();
this.compactionQueuedMessages = [];
this.streamingComponent = undefined;
this.streamingMessage = undefined;
this.pendingTools.clear();
this.chatContainer.addChild(new Spacer(1)); this.chatContainer.addChild(new Spacer(1));
this.chatContainer.addChild(new Text(`${theme.fg("accent", "✓ New session started")}`, 1, 1)); this.chatContainer.addChild(new Text(`${theme.fg("accent", "✓ New session started")}`, 1, 1));
this.ui.requestRender(); this.ui.requestRender();
@@ -4511,7 +4503,7 @@ export class InteractiveMode {
type: "user_bash", type: "user_bash",
command, command,
excludeFromContext, excludeFromContext,
cwd: process.cwd(), cwd: this.sessionManager.getCwd(),
}) })
: undefined; : undefined;

View File

@@ -7,7 +7,7 @@
*/ */
import type { AssistantMessage, ImageContent } from "@mariozechner/pi-ai"; import type { AssistantMessage, ImageContent } from "@mariozechner/pi-ai";
import type { AgentSession } from "../core/agent-session.js"; import type { AgentSessionRuntimeHost } from "../core/agent-session-runtime.js";
import { flushRawStdout, writeRawStdout } from "../core/output-guard.js"; import { flushRawStdout, writeRawStdout } from "../core/output-guard.js";
/** /**
@@ -28,44 +28,46 @@ export interface PrintModeOptions {
* Run in print (single-shot) mode. * Run in print (single-shot) mode.
* Sends prompts to the agent and outputs the result. * Sends prompts to the agent and outputs the result.
*/ */
export async function runPrintMode(session: AgentSession, options: PrintModeOptions): Promise<number> { export async function runPrintMode(runtimeHost: AgentSessionRuntimeHost, options: PrintModeOptions): Promise<number> {
const { mode, messages = [], initialMessage, initialImages } = options; const { mode, messages = [], initialMessage, initialImages } = options;
let exitCode = 0; let exitCode = 0;
let session = runtimeHost.session;
let unsubscribe: (() => void) | undefined;
try { const rebindSession = async (): Promise<void> => {
if (mode === "json") { session = runtimeHost.session;
const header = session.sessionManager.getHeader();
if (header) {
writeRawStdout(`${JSON.stringify(header)}\n`);
}
}
// Set up extensions for print mode (no UI)
await session.bindExtensions({ await session.bindExtensions({
commandContextActions: { commandContextActions: {
waitForIdle: () => session.agent.waitForIdle(), waitForIdle: () => session.agent.waitForIdle(),
newSession: async (options) => { newSession: async (newSessionOptions) => {
const success = await session.newSession({ parentSession: options?.parentSession }); const result = await runtimeHost.newSession(newSessionOptions);
if (success && options?.setup) { if (!result.cancelled) {
await options.setup(session.sessionManager); await rebindSession();
} }
return { cancelled: !success }; return result;
}, },
fork: async (entryId) => { fork: async (entryId) => {
const result = await session.fork(entryId); const result = await runtimeHost.fork(entryId);
if (!result.cancelled) {
await rebindSession();
}
return { cancelled: result.cancelled }; return { cancelled: result.cancelled };
}, },
navigateTree: async (targetId, options) => { navigateTree: async (targetId, navigateOptions) => {
const result = await session.navigateTree(targetId, { const result = await session.navigateTree(targetId, {
summarize: options?.summarize, summarize: navigateOptions?.summarize,
customInstructions: options?.customInstructions, customInstructions: navigateOptions?.customInstructions,
replaceInstructions: options?.replaceInstructions, replaceInstructions: navigateOptions?.replaceInstructions,
label: options?.label, label: navigateOptions?.label,
}); });
return { cancelled: result.cancelled }; return { cancelled: result.cancelled };
}, },
switchSession: async (sessionPath) => { switchSession: async (sessionPath) => {
const success = await session.switchSession(sessionPath); const result = await runtimeHost.switchSession(sessionPath);
return { cancelled: !success }; if (!result.cancelled) {
await rebindSession();
}
return result;
}, },
reload: async () => { reload: async () => {
await session.reload(); await session.reload();
@@ -76,38 +78,42 @@ export async function runPrintMode(session: AgentSession, options: PrintModeOpti
}, },
}); });
// Always subscribe to enable session persistence via _handleAgentEvent unsubscribe?.();
session.subscribe((event) => { unsubscribe = session.subscribe((event) => {
// In JSON mode, output all events
if (mode === "json") { if (mode === "json") {
writeRawStdout(`${JSON.stringify(event)}\n`); writeRawStdout(`${JSON.stringify(event)}\n`);
} }
}); });
};
try {
if (mode === "json") {
const header = session.sessionManager.getHeader();
if (header) {
writeRawStdout(`${JSON.stringify(header)}\n`);
}
}
await rebindSession();
// Send initial message with attachments
if (initialMessage) { if (initialMessage) {
await session.prompt(initialMessage, { images: initialImages }); await session.prompt(initialMessage, { images: initialImages });
} }
// Send remaining messages
for (const message of messages) { for (const message of messages) {
await session.prompt(message); await session.prompt(message);
} }
// In text mode, output final response
if (mode === "text") { if (mode === "text") {
const state = session.state; const state = session.state;
const lastMessage = state.messages[state.messages.length - 1]; const lastMessage = state.messages[state.messages.length - 1];
if (lastMessage?.role === "assistant") { if (lastMessage?.role === "assistant") {
const assistantMsg = lastMessage as AssistantMessage; const assistantMsg = lastMessage as AssistantMessage;
// Check for error/aborted
if (assistantMsg.stopReason === "error" || assistantMsg.stopReason === "aborted") { if (assistantMsg.stopReason === "error" || assistantMsg.stopReason === "aborted") {
console.error(assistantMsg.errorMessage || `Request ${assistantMsg.stopReason}`); console.error(assistantMsg.errorMessage || `Request ${assistantMsg.stopReason}`);
exitCode = 1; exitCode = 1;
} else { } else {
// Output text content
for (const content of assistantMsg.content) { for (const content of assistantMsg.content) {
if (content.type === "text") { if (content.type === "text") {
writeRawStdout(`${content.text}\n`); writeRawStdout(`${content.text}\n`);
@@ -119,13 +125,8 @@ export async function runPrintMode(session: AgentSession, options: PrintModeOpti
return exitCode; return exitCode;
} finally { } finally {
const extensionRunner = session.extensionRunner; unsubscribe?.();
if (extensionRunner?.hasHandlers("session_shutdown")) { await runtimeHost.dispose();
await extensionRunner.emit({ type: "session_shutdown" });
}
// Ensure stdout is fully flushed before returning
// This prevents race conditions where the process exits before all output is written
await flushRawStdout(); await flushRawStdout();
} }
} }

View File

@@ -12,7 +12,7 @@
*/ */
import * as crypto from "node:crypto"; import * as crypto from "node:crypto";
import type { AgentSession } from "../../core/agent-session.js"; import type { AgentSessionRuntimeHost } from "../../core/agent-session-runtime.js";
import type { import type {
ExtensionUIContext, ExtensionUIContext,
ExtensionUIDialogOptions, ExtensionUIDialogOptions,
@@ -43,8 +43,10 @@ export type {
* Run in RPC mode. * Run in RPC mode.
* Listens for JSON commands on stdin, outputs events and responses on stdout. * Listens for JSON commands on stdin, outputs events and responses on stdout.
*/ */
export async function runRpcMode(session: AgentSession): Promise<never> { export async function runRpcMode(runtimeHost: AgentSessionRuntimeHost): Promise<never> {
takeOverStdout(); takeOverStdout();
let session = runtimeHost.session;
let unsubscribe: (() => void) | undefined;
const output = (obj: RpcResponse | RpcExtensionUIRequest | object) => { const output = (obj: RpcResponse | RpcExtensionUIRequest | object) => {
writeRawStdout(serializeJsonLine(obj)); writeRawStdout(serializeJsonLine(obj));
@@ -280,49 +282,61 @@ export async function runRpcMode(session: AgentSession): Promise<never> {
}, },
}); });
// Set up extensions with RPC-based UI context const rebindSession = async (): Promise<void> => {
await session.bindExtensions({ session = runtimeHost.session;
uiContext: createExtensionUIContext(), await session.bindExtensions({
commandContextActions: { uiContext: createExtensionUIContext(),
waitForIdle: () => session.agent.waitForIdle(), commandContextActions: {
newSession: async (options) => { waitForIdle: () => session.agent.waitForIdle(),
// Delegate to AgentSession (handles setup + agent state sync) newSession: async (options) => {
const success = await session.newSession(options); const result = await runtimeHost.newSession(options);
return { cancelled: !success }; if (!result.cancelled) {
await rebindSession();
}
return result;
},
fork: async (entryId) => {
const result = await runtimeHost.fork(entryId);
if (!result.cancelled) {
await rebindSession();
}
return { cancelled: result.cancelled };
},
navigateTree: async (targetId, options) => {
const result = await session.navigateTree(targetId, {
summarize: options?.summarize,
customInstructions: options?.customInstructions,
replaceInstructions: options?.replaceInstructions,
label: options?.label,
});
return { cancelled: result.cancelled };
},
switchSession: async (sessionPath) => {
const result = await runtimeHost.switchSession(sessionPath);
if (!result.cancelled) {
await rebindSession();
}
return result;
},
reload: async () => {
await session.reload();
},
}, },
fork: async (entryId) => { shutdownHandler: () => {
const result = await session.fork(entryId); shutdownRequested = true;
return { cancelled: result.cancelled };
}, },
navigateTree: async (targetId, options) => { onError: (err) => {
const result = await session.navigateTree(targetId, { output({ type: "extension_error", extensionPath: err.extensionPath, event: err.event, error: err.error });
summarize: options?.summarize,
customInstructions: options?.customInstructions,
replaceInstructions: options?.replaceInstructions,
label: options?.label,
});
return { cancelled: result.cancelled };
}, },
switchSession: async (sessionPath) => { });
const success = await session.switchSession(sessionPath);
return { cancelled: !success };
},
reload: async () => {
await session.reload();
},
},
shutdownHandler: () => {
shutdownRequested = true;
},
onError: (err) => {
output({ type: "extension_error", extensionPath: err.extensionPath, event: err.event, error: err.error });
},
});
// Output all agent events as JSON unsubscribe?.();
session.subscribe((event) => { unsubscribe = session.subscribe((event) => {
output(event); output(event);
}); });
};
await rebindSession();
// Handle a single command // Handle a single command
const handleCommand = async (command: RpcCommand): Promise<RpcResponse> => { const handleCommand = async (command: RpcCommand): Promise<RpcResponse> => {
@@ -364,8 +378,11 @@ export async function runRpcMode(session: AgentSession): Promise<never> {
case "new_session": { case "new_session": {
const options = command.parentSession ? { parentSession: command.parentSession } : undefined; const options = command.parentSession ? { parentSession: command.parentSession } : undefined;
const cancelled = !(await session.newSession(options)); const result = await runtimeHost.newSession(options);
return success(id, "new_session", { cancelled }); if (!result.cancelled) {
await rebindSession();
}
return success(id, "new_session", result);
} }
// ================================================================= // =================================================================
@@ -505,12 +522,18 @@ export async function runRpcMode(session: AgentSession): Promise<never> {
} }
case "switch_session": { case "switch_session": {
const cancelled = !(await session.switchSession(command.sessionPath)); const result = await runtimeHost.switchSession(command.sessionPath);
return success(id, "switch_session", { cancelled }); if (!result.cancelled) {
await rebindSession();
}
return success(id, "switch_session", result);
} }
case "fork": { case "fork": {
const result = await session.fork(command.entryId); const result = await runtimeHost.fork(command.entryId);
if (!result.cancelled) {
await rebindSession();
}
return success(id, "fork", { text: result.selectedText, cancelled: result.cancelled }); return success(id, "fork", { text: result.selectedText, cancelled: result.cancelled });
} }
@@ -592,11 +615,8 @@ export async function runRpcMode(session: AgentSession): Promise<never> {
let detachInput = () => {}; let detachInput = () => {};
async function shutdown(): Promise<never> { async function shutdown(): Promise<never> {
const currentRunner = session.extensionRunner; unsubscribe?.();
if (currentRunner?.hasHandlers("session_shutdown")) { await runtimeHost.dispose();
await currentRunner.emit({ type: "session_shutdown" });
}
detachInput(); detachInput();
process.stdin.pause(); process.stdin.pause();
process.exit(0); process.exit(0);

View File

@@ -10,126 +10,112 @@
import { existsSync, mkdirSync, rmSync } from "node:fs"; import { existsSync, mkdirSync, rmSync } from "node:fs";
import { tmpdir } from "node:os"; import { tmpdir } from "node:os";
import { join } from "node:path"; import { join } from "node:path";
import { Agent } from "@mariozechner/pi-agent-core";
import { getModel } from "@mariozechner/pi-ai"; import { getModel } from "@mariozechner/pi-ai";
import { afterEach, beforeEach, describe, expect, it } from "vitest"; import { afterEach, beforeEach, describe, expect, it } from "vitest";
import { AgentSession } from "../src/core/agent-session.js"; import type { AgentSession } from "../src/core/agent-session.js";
import {
type AgentSessionRuntimeHost,
createAgentSessionRuntime,
AgentSessionRuntimeHost as RuntimeHost,
} from "../src/core/agent-session-runtime.js";
import { AuthStorage } from "../src/core/auth-storage.js"; import { AuthStorage } from "../src/core/auth-storage.js";
import { ModelRegistry } from "../src/core/model-registry.js";
import { SessionManager } from "../src/core/session-manager.js"; import { SessionManager } from "../src/core/session-manager.js";
import { SettingsManager } from "../src/core/settings-manager.js";
import { codingTools } from "../src/core/tools/index.js"; import { codingTools } from "../src/core/tools/index.js";
import { API_KEY, createTestResourceLoader } from "./utilities.js"; import { API_KEY } from "./utilities.js";
describe.skipIf(!API_KEY)("AgentSession forking", () => { describe.skipIf(!API_KEY)("AgentSession forking", () => {
let session: AgentSession; let session: AgentSession;
let runtimeHost: AgentSessionRuntimeHost;
let tempDir: string; let tempDir: string;
let sessionManager: SessionManager; let sessionManager: SessionManager;
beforeEach(() => { beforeEach(() => {
// Create temp directory for session files
tempDir = join(tmpdir(), `pi-branching-test-${Date.now()}`); tempDir = join(tmpdir(), `pi-branching-test-${Date.now()}`);
mkdirSync(tempDir, { recursive: true }); mkdirSync(tempDir, { recursive: true });
}); });
afterEach(async () => { afterEach(async () => {
if (session) { if (runtimeHost) {
session.dispose(); await runtimeHost.dispose();
} }
process.chdir(tmpdir());
if (tempDir && existsSync(tempDir)) { if (tempDir && existsSync(tempDir)) {
rmSync(tempDir, { recursive: true }); rmSync(tempDir, { recursive: true });
} }
}); });
function createSession(noSession: boolean = false) { async function createSession(noSession: boolean = false) {
const model = getModel("anthropic", "claude-sonnet-4-5")!; const model = getModel("anthropic", "claude-sonnet-4-5")!;
const agent = new Agent({ sessionManager = noSession ? SessionManager.inMemory(tempDir) : SessionManager.create(tempDir);
getApiKey: () => API_KEY,
initialState: {
model,
systemPrompt: "You are a helpful assistant. Be extremely concise, reply with just a few words.",
tools: codingTools,
},
});
sessionManager = noSession ? SessionManager.inMemory() : SessionManager.create(tempDir);
const settingsManager = SettingsManager.create(tempDir, tempDir);
const authStorage = AuthStorage.create(join(tempDir, "auth.json")); const authStorage = AuthStorage.create(join(tempDir, "auth.json"));
const modelRegistry = ModelRegistry.create(authStorage, tempDir); authStorage.setRuntimeApiKey("anthropic", API_KEY!);
session = new AgentSession({ const bootstrap = {
agent, agentDir: tempDir,
sessionManager, authStorage,
settingsManager, model,
tools: codingTools,
resourceLoader: {
noExtensions: true,
noSkills: true,
noPromptTemplates: true,
noThemes: true,
},
};
const runtime = await createAgentSessionRuntime(bootstrap, {
cwd: tempDir, cwd: tempDir,
modelRegistry, sessionManager,
resourceLoader: createTestResourceLoader(),
}); });
runtimeHost = new RuntimeHost(bootstrap, runtime);
// Must subscribe to enable session persistence session = runtimeHost.session;
session.subscribe(() => {}); session.subscribe(() => {});
return session; return session;
} }
it("should allow forking from single message", async () => { it("should allow forking from single message", async () => {
createSession(); await createSession();
// Send one message
await session.prompt("Say hello"); await session.prompt("Say hello");
await session.agent.waitForIdle(); await session.agent.waitForIdle();
// Should have exactly 1 user message available for forking
const userMessages = session.getUserMessagesForForking(); const userMessages = session.getUserMessagesForForking();
expect(userMessages.length).toBe(1); expect(userMessages.length).toBe(1);
expect(userMessages[0].text).toBe("Say hello"); expect(userMessages[0].text).toBe("Say hello");
// Fork from the first message const result = await runtimeHost.fork(userMessages[0].entryId);
const result = await session.fork(userMessages[0].entryId);
expect(result.selectedText).toBe("Say hello");
expect(result.cancelled).toBe(false); expect(result.cancelled).toBe(false);
session = runtimeHost.session;
expect(result.selectedText).toBe("Say hello");
// After forking, conversation should be empty (forked before the first message)
expect(session.messages.length).toBe(0); expect(session.messages.length).toBe(0);
// Session file path should be set, but file is created lazily after first assistant message
expect(session.sessionFile).not.toBeNull(); expect(session.sessionFile).not.toBeNull();
expect(existsSync(session.sessionFile!)).toBe(false); expect(existsSync(session.sessionFile!)).toBe(false);
}); });
it("should support in-memory forking in --no-session mode", async () => { it("should support in-memory forking in --no-session mode", async () => {
createSession(true); await createSession(true);
// Verify sessions are disabled
expect(session.sessionFile).toBeUndefined(); expect(session.sessionFile).toBeUndefined();
// Send one message
await session.prompt("Say hi"); await session.prompt("Say hi");
await session.agent.waitForIdle(); await session.agent.waitForIdle();
// Should have 1 user message
const userMessages = session.getUserMessagesForForking(); const userMessages = session.getUserMessagesForForking();
expect(userMessages.length).toBe(1); expect(userMessages.length).toBe(1);
// Verify we have messages before forking
expect(session.messages.length).toBeGreaterThan(0); expect(session.messages.length).toBeGreaterThan(0);
// Fork from the first message const result = await runtimeHost.fork(userMessages[0].entryId);
const result = await session.fork(userMessages[0].entryId);
expect(result.selectedText).toBe("Say hi");
expect(result.cancelled).toBe(false); expect(result.cancelled).toBe(false);
session = runtimeHost.session;
expect(result.selectedText).toBe("Say hi");
// After forking, conversation should be empty
expect(session.messages.length).toBe(0); expect(session.messages.length).toBe(0);
// Session file should still be undefined (no file created)
expect(session.sessionFile).toBeUndefined(); expect(session.sessionFile).toBeUndefined();
}); });
it("should fork from middle of conversation", async () => { it("should fork from middle of conversation", async () => {
createSession(); await createSession();
// Send multiple messages
await session.prompt("Say one"); await session.prompt("Say one");
await session.agent.waitForIdle(); await session.agent.waitForIdle();
@@ -139,18 +125,17 @@ describe.skipIf(!API_KEY)("AgentSession forking", () => {
await session.prompt("Say three"); await session.prompt("Say three");
await session.agent.waitForIdle(); await session.agent.waitForIdle();
// Should have 3 user messages
const userMessages = session.getUserMessagesForForking(); const userMessages = session.getUserMessagesForForking();
expect(userMessages.length).toBe(3); expect(userMessages.length).toBe(3);
// Fork from second message (keeps first message + response)
const secondMessage = userMessages[1]; const secondMessage = userMessages[1];
const result = await session.fork(secondMessage.entryId); const result = await runtimeHost.fork(secondMessage.entryId);
expect(result.cancelled).toBe(false);
session = runtimeHost.session;
expect(result.selectedText).toBe("Say two"); expect(result.selectedText).toBe("Say two");
// After forking, should have first user message + assistant response
expect(session.messages.length).toBe(2); expect(session.messages.length).toBe(2);
expect(session.messages[0].role).toBe("user"); expect(session.messages[0].role).toBe("user");
expect(session.messages[1].role).toBe("assistant"); expect(session.messages[1].role).toBe("assistant");
}); }, 60000);
}); });

View File

@@ -0,0 +1,172 @@
import { existsSync, mkdirSync, rmSync } from "node:fs";
import { tmpdir } from "node:os";
import { join } from "node:path";
import { fauxAssistantMessage, registerFauxProvider } from "@mariozechner/pi-ai";
import { afterEach, describe, expect, it } from "vitest";
import {
type AgentSessionRuntimeBootstrap,
AgentSessionRuntimeHost,
createAgentSessionRuntime,
} from "../src/core/agent-session-runtime.js";
import { AuthStorage } from "../src/core/auth-storage.js";
import { SessionManager } from "../src/core/session-manager.js";
import type {
ExtensionFactory,
SessionBeforeForkEvent,
SessionBeforeSwitchEvent,
SessionStartEvent,
} from "../src/index.js";
type RecordedSessionEvent = SessionBeforeSwitchEvent | SessionBeforeForkEvent | SessionStartEvent;
describe("AgentSessionRuntimeHost session lifecycle events", () => {
const cleanups: Array<() => Promise<void> | void> = [];
afterEach(async () => {
while (cleanups.length > 0) {
await cleanups.pop()?.();
}
});
async function createRuntimeHost(extensionFactory: ExtensionFactory) {
const tempDir = join(tmpdir(), `pi-runtime-events-${Date.now()}-${Math.random().toString(36).slice(2)}`);
mkdirSync(tempDir, { recursive: true });
const faux = registerFauxProvider();
faux.setResponses([fauxAssistantMessage("one"), fauxAssistantMessage("two"), fauxAssistantMessage("three")]);
const authStorage = AuthStorage.inMemory();
authStorage.setRuntimeApiKey(faux.getModel().provider, "faux-key");
const bootstrap: AgentSessionRuntimeBootstrap = {
agentDir: tempDir,
authStorage,
model: faux.getModel(),
resourceLoader: {
extensionFactories: [extensionFactory],
noSkills: true,
noPromptTemplates: true,
noThemes: true,
},
};
const runtime = await createAgentSessionRuntime(bootstrap, {
cwd: tempDir,
sessionManager: SessionManager.create(tempDir),
});
const runtimeHost = new AgentSessionRuntimeHost(bootstrap, runtime);
await runtimeHost.session.bindExtensions({});
cleanups.push(async () => {
await runtimeHost.dispose();
faux.unregister();
process.chdir(tmpdir());
if (existsSync(tempDir)) {
rmSync(tempDir, { recursive: true, force: true });
}
});
return { runtimeHost, faux };
}
it("emits session_before_switch and session_start for new and resume flows", async () => {
const events: RecordedSessionEvent[] = [];
const { runtimeHost } = await createRuntimeHost((pi) => {
pi.on("session_before_switch", (event) => {
events.push(event);
});
pi.on("session_start", (event) => {
events.push(event);
});
});
expect(events).toEqual([{ type: "session_start", reason: "startup" }]);
events.length = 0;
await runtimeHost.session.prompt("hello");
const originalSessionFile = runtimeHost.session.sessionFile;
expect(originalSessionFile).toBeTruthy();
const newSessionResult = await runtimeHost.newSession();
expect(newSessionResult.cancelled).toBe(false);
await runtimeHost.session.bindExtensions({});
expect(events).toEqual([
{ type: "session_before_switch", reason: "new", targetSessionFile: undefined },
{ type: "session_start", reason: "new", previousSessionFile: originalSessionFile },
]);
events.length = 0;
const secondSessionFile = runtimeHost.session.sessionFile;
expect(secondSessionFile).toBeTruthy();
const switchResult = await runtimeHost.switchSession(originalSessionFile!);
expect(switchResult.cancelled).toBe(false);
await runtimeHost.session.bindExtensions({});
expect(events).toEqual([
{ type: "session_before_switch", reason: "resume", targetSessionFile: originalSessionFile },
{ type: "session_start", reason: "resume", previousSessionFile: secondSessionFile },
]);
});
it("honors session_before_switch cancellation", async () => {
const events: RecordedSessionEvent[] = [];
const { runtimeHost } = await createRuntimeHost((pi) => {
pi.on("session_before_switch", (event) => {
events.push(event);
return { cancel: true };
});
pi.on("session_start", (event) => {
events.push(event);
});
});
expect(events).toEqual([{ type: "session_start", reason: "startup" }]);
events.length = 0;
await runtimeHost.session.prompt("hello");
const originalSessionFile = runtimeHost.session.sessionFile;
const result = await runtimeHost.newSession();
expect(result.cancelled).toBe(true);
expect(runtimeHost.session.sessionFile).toBe(originalSessionFile);
expect(events).toEqual([{ type: "session_before_switch", reason: "new", targetSessionFile: undefined }]);
});
it("emits session_before_fork and session_start and honors cancellation", async () => {
const events: RecordedSessionEvent[] = [];
let cancelNextFork = false;
const { runtimeHost } = await createRuntimeHost((pi) => {
pi.on("session_before_fork", (event) => {
events.push(event);
if (cancelNextFork) {
cancelNextFork = false;
return { cancel: true };
}
});
pi.on("session_start", (event) => {
events.push(event);
});
});
expect(events).toEqual([{ type: "session_start", reason: "startup" }]);
events.length = 0;
await runtimeHost.session.prompt("hello");
const userMessage = runtimeHost.session.getUserMessagesForForking()[0];
const previousSessionFile = runtimeHost.session.sessionFile;
const successResult = await runtimeHost.fork(userMessage.entryId);
expect(successResult.cancelled).toBe(false);
expect(successResult.selectedText).toBe("hello");
await runtimeHost.session.bindExtensions({});
expect(events).toEqual([
{ type: "session_before_fork", entryId: userMessage.entryId },
{ type: "session_start", reason: "fork", previousSessionFile },
]);
events.length = 0;
cancelNextFork = true;
const cancelResult = await runtimeHost.fork(userMessage.entryId);
expect(cancelResult).toEqual({ cancelled: true });
expect(events).toEqual([{ type: "session_before_fork", entryId: userMessage.entryId }]);
});
});

View File

@@ -48,6 +48,7 @@ function createSession(options: {
sessionManager: { sessionManager: {
getEntries: () => entries, getEntries: () => entries,
getSessionName: () => options.sessionName, getSessionName: () => options.sessionName,
getCwd: () => "/tmp/project",
}, },
getContextUsage: () => ({ contextWindow: 200_000, percent: 12.3 }), getContextUsage: () => ({ contextWindow: 200_000, percent: 12.3 }),
modelRegistry: { modelRegistry: {

View File

@@ -17,11 +17,15 @@ type FakeSession = {
bindExtensions: ReturnType<typeof vi.fn>; bindExtensions: ReturnType<typeof vi.fn>;
subscribe: ReturnType<typeof vi.fn>; subscribe: ReturnType<typeof vi.fn>;
prompt: ReturnType<typeof vi.fn>; prompt: ReturnType<typeof vi.fn>;
reload: ReturnType<typeof vi.fn>;
};
type FakeRuntimeHost = {
session: FakeSession;
newSession: ReturnType<typeof vi.fn>; newSession: ReturnType<typeof vi.fn>;
fork: ReturnType<typeof vi.fn>; fork: ReturnType<typeof vi.fn>;
navigateTree: ReturnType<typeof vi.fn>;
switchSession: ReturnType<typeof vi.fn>; switchSession: ReturnType<typeof vi.fn>;
reload: ReturnType<typeof vi.fn>; dispose: ReturnType<typeof vi.fn>;
}; };
function createAssistantMessage(options?: { function createAssistantMessage(options?: {
@@ -49,7 +53,7 @@ function createAssistantMessage(options?: {
}; };
} }
function createSession(assistantMessage: AssistantMessage): FakeSession { function createRuntimeHost(assistantMessage: AssistantMessage): FakeRuntimeHost {
const extensionRunner: FakeExtensionRunner = { const extensionRunner: FakeExtensionRunner = {
hasHandlers: (eventType: string) => eventType === "session_shutdown", hasHandlers: (eventType: string) => eventType === "session_shutdown",
emit: vi.fn(async () => {}), emit: vi.fn(async () => {}),
@@ -57,7 +61,7 @@ function createSession(assistantMessage: AssistantMessage): FakeSession {
const state = { messages: [assistantMessage] }; const state = { messages: [assistantMessage] };
return { const session: FakeSession = {
sessionManager: { getHeader: () => undefined }, sessionManager: { getHeader: () => undefined },
agent: { waitForIdle: async () => {} }, agent: { waitForIdle: async () => {} },
state, state,
@@ -65,12 +69,18 @@ function createSession(assistantMessage: AssistantMessage): FakeSession {
bindExtensions: vi.fn(async () => {}), bindExtensions: vi.fn(async () => {}),
subscribe: vi.fn(() => () => {}), subscribe: vi.fn(() => () => {}),
prompt: vi.fn(async () => {}), prompt: vi.fn(async () => {}),
newSession: vi.fn(async () => true),
fork: vi.fn(async () => ({ cancelled: false })),
navigateTree: vi.fn(async () => ({ cancelled: false })),
switchSession: vi.fn(async () => true),
reload: vi.fn(async () => {}), reload: vi.fn(async () => {}),
}; };
return {
session,
newSession: vi.fn(async () => undefined),
fork: vi.fn(async () => ({ selectedText: "" })),
switchSession: vi.fn(async () => undefined),
dispose: vi.fn(async () => {
await session.extensionRunner.emit({ type: "session_shutdown" });
}),
};
} }
afterEach(() => { afterEach(() => {
@@ -79,10 +89,11 @@ afterEach(() => {
describe("runPrintMode", () => { describe("runPrintMode", () => {
it("emits session_shutdown in text mode", async () => { it("emits session_shutdown in text mode", async () => {
const session = createSession(createAssistantMessage({ text: "done" })); const runtimeHost = createRuntimeHost(createAssistantMessage({ text: "done" }));
const { session } = runtimeHost;
const images: ImageContent[] = [{ type: "image", mimeType: "image/png", data: "abc" }]; const images: ImageContent[] = [{ type: "image", mimeType: "image/png", data: "abc" }];
const exitCode = await runPrintMode(session as unknown as Parameters<typeof runPrintMode>[0], { const exitCode = await runPrintMode(runtimeHost as unknown as Parameters<typeof runPrintMode>[0], {
mode: "text", mode: "text",
initialMessage: "Say done", initialMessage: "Say done",
initialImages: images, initialImages: images,
@@ -95,9 +106,10 @@ describe("runPrintMode", () => {
}); });
it("emits session_shutdown in json mode", async () => { it("emits session_shutdown in json mode", async () => {
const session = createSession(createAssistantMessage({ text: "done" })); const runtimeHost = createRuntimeHost(createAssistantMessage({ text: "done" }));
const { session } = runtimeHost;
const exitCode = await runPrintMode(session as unknown as Parameters<typeof runPrintMode>[0], { const exitCode = await runPrintMode(runtimeHost as unknown as Parameters<typeof runPrintMode>[0], {
mode: "json", mode: "json",
messages: ["hello"], messages: ["hello"],
}); });
@@ -109,10 +121,13 @@ describe("runPrintMode", () => {
}); });
it("emits session_shutdown and returns non-zero on assistant error", async () => { it("emits session_shutdown and returns non-zero on assistant error", async () => {
const session = createSession(createAssistantMessage({ stopReason: "error", errorMessage: "provider failure" })); const runtimeHost = createRuntimeHost(
createAssistantMessage({ stopReason: "error", errorMessage: "provider failure" }),
);
const { session } = runtimeHost;
const errorSpy = vi.spyOn(console, "error").mockImplementation(() => {}); const errorSpy = vi.spyOn(console, "error").mockImplementation(() => {});
const exitCode = await runPrintMode(session as unknown as Parameters<typeof runPrintMode>[0], { const exitCode = await runPrintMode(runtimeHost as unknown as Parameters<typeof runPrintMode>[0], {
mode: "text", mode: "text",
}); });