fix: align embedded runner with session API changes

This commit is contained in:
Peter Steinberger
2026-02-01 15:06:42 -08:00
parent bcbb447357
commit 3367b2aa27
4 changed files with 55 additions and 15 deletions

View File

@@ -18,7 +18,7 @@ import { ensureAuthProfileStore, saveAuthProfileStore } from "./store.js";
const OAUTH_PROVIDER_IDS = new Set(getOAuthProviders().map((provider) => provider.id));
const resolveOAuthProvider = (provider: string): OAuthProvider | null =>
OAUTH_PROVIDER_IDS.has(provider as OAuthProvider) ? (provider as OAuthProvider) : null;
OAUTH_PROVIDER_IDS.has(provider) ? provider : null;
function buildOAuthApiKey(provider: string, credentials: OAuthCredentials): string {
const needsProjectId = provider === "google-gemini-cli" || provider === "google-antigravity";

View File

@@ -1,6 +1,7 @@
import {
createAgentSession,
estimateTokens,
DefaultResourceLoader,
SessionManager,
SettingsManager,
} from "@mariozechner/pi-coding-agent";
@@ -64,7 +65,11 @@ import { log } from "./logger.js";
import { buildModelAliasLines, resolveModel } from "./model.js";
import { buildEmbeddedSandboxInfo } from "./sandbox-info.js";
import { prewarmSessionFile, trackSessionManagerAccess } from "./session-manager-cache.js";
import { buildEmbeddedSystemPrompt, createSystemPromptOverride } from "./system-prompt.js";
import {
applySystemPromptOverrideToSession,
buildEmbeddedSystemPrompt,
createSystemPromptOverride,
} from "./system-prompt.js";
import { splitSdkTools } from "./tool-split.js";
import { describeUnknownError, mapThinkingLevel, resolveExecToolDefaults } from "./utils.js";
@@ -347,7 +352,7 @@ export async function compactEmbeddedPiSessionDirect(
userTimeFormat,
contextFiles,
});
const systemPrompt = createSystemPromptOverride(appendPrompt);
const systemPromptOverride = createSystemPromptOverride(appendPrompt);
const sessionLock = await acquireSessionWriteLock({
sessionFile: params.sessionFile,
@@ -383,6 +388,13 @@ export async function compactEmbeddedPiSessionDirect(
sandboxEnabled: !!sandbox?.enabled,
});
const resourceLoader = new DefaultResourceLoader({
cwd: resolvedWorkspace,
agentDir,
settingsManager,
additionalExtensionPaths,
});
await resourceLoader.reload();
const { session } = await createAgentSession({
cwd: resolvedWorkspace,
agentDir,
@@ -394,11 +406,9 @@ export async function compactEmbeddedPiSessionDirect(
customTools,
sessionManager,
settingsManager,
systemPrompt,
additionalExtensionPaths,
skills: [],
contextFiles: [],
resourceLoader,
});
applySystemPromptOverrideToSession(session, systemPromptOverride);
try {
const prior = await sanitizeSessionHistory({

View File

@@ -1,7 +1,12 @@
import type { AgentMessage } from "@mariozechner/pi-agent-core";
import type { ImageContent } from "@mariozechner/pi-ai";
import { streamSimple } from "@mariozechner/pi-ai";
import { createAgentSession, SessionManager, SettingsManager } from "@mariozechner/pi-coding-agent";
import {
createAgentSession,
DefaultResourceLoader,
SessionManager,
SettingsManager,
} from "@mariozechner/pi-coding-agent";
import fs from "node:fs/promises";
import os from "node:os";
import type { EmbeddedRunAttemptParams, EmbeddedRunAttemptResult } from "./types.js";
@@ -78,7 +83,11 @@ import {
import { buildEmbeddedSandboxInfo } from "../sandbox-info.js";
import { prewarmSessionFile, trackSessionManagerAccess } from "../session-manager-cache.js";
import { prepareSessionManagerForRun } from "../session-manager-init.js";
import { buildEmbeddedSystemPrompt, createSystemPromptOverride } from "../system-prompt.js";
import {
applySystemPromptOverrideToSession,
buildEmbeddedSystemPrompt,
createSystemPromptOverride,
} from "../system-prompt.js";
import { splitSdkTools } from "../tool-split.js";
import { describeUnknownError, mapThinkingLevel } from "../utils.js";
import { detectAndLoadPromptImages } from "./images.js";
@@ -385,7 +394,8 @@ export async function runEmbeddedAttempt(
skillsPrompt,
tools,
});
const systemPrompt = createSystemPromptOverride(appendPrompt);
const systemPromptOverride = createSystemPromptOverride(appendPrompt);
const systemPromptText = systemPromptOverride();
const sessionLock = await acquireSessionWriteLock({
sessionFile: params.sessionFile,
@@ -457,6 +467,13 @@ export async function runEmbeddedAttempt(
const allCustomTools = [...customTools, ...clientToolDefs];
const resourceLoader = new DefaultResourceLoader({
cwd: resolvedWorkspace,
agentDir,
settingsManager,
additionalExtensionPaths,
});
await resourceLoader.reload();
({ session } = await createAgentSession({
cwd: resolvedWorkspace,
agentDir,
@@ -468,11 +485,9 @@ export async function runEmbeddedAttempt(
customTools: allCustomTools,
sessionManager,
settingsManager,
systemPrompt,
additionalExtensionPaths,
skills: [],
contextFiles: [],
resourceLoader,
}));
applySystemPromptOverrideToSession(session, systemPromptOverride);
if (!session) {
throw new Error("Embedded agent session missing");
}
@@ -513,7 +528,7 @@ export async function runEmbeddedAttempt(
if (cacheTrace) {
cacheTrace.recordStage("session:loaded", {
messages: activeSession.messages,
system: systemPrompt,
system: systemPromptText,
note: "after session create",
});
activeSession.agent.streamFn = cacheTrace.wrapStreamFn(activeSession.agent.streamFn);

View File

@@ -1,4 +1,5 @@
import type { AgentTool } from "@mariozechner/pi-agent-core";
import type { AgentSession } from "@mariozechner/pi-coding-agent";
import type { ResolvedTimeFormat } from "../date-time.js";
import type { EmbeddedContextFile } from "../pi-embedded-helpers.js";
import type { EmbeddedSandboxInfo } from "./types.js";
@@ -79,3 +80,17 @@ export function createSystemPromptOverride(
const trimmed = systemPrompt.trim();
return () => trimmed;
}
export function applySystemPromptOverrideToSession(
session: AgentSession,
override: (defaultPrompt?: string) => string,
) {
const prompt = override().trim();
session.agent.setSystemPrompt(prompt);
const mutableSession = session as unknown as {
_baseSystemPrompt?: string;
_rebuildSystemPrompt?: (toolNames: string[]) => string;
};
mutableSession._baseSystemPrompt = prompt;
mutableSession._rebuildSystemPrompt = () => prompt;
}