diff --git a/app/src/client/lib/desk.ts b/app/src/client/lib/desk.ts index 0073075..1d929ce 100644 --- a/app/src/client/lib/desk.ts +++ b/app/src/client/lib/desk.ts @@ -7,6 +7,7 @@ export type DeskState = { settings: PiSettings | null; appErrors: Partial>; authUrls?: Partial>; + gcalReady?: boolean; }; export type McpServersSnapshot = { diff --git a/app/src/client/pages/AppsPage.tsx b/app/src/client/pages/AppsPage.tsx index e7e7ff5..3599ed6 100644 --- a/app/src/client/pages/AppsPage.tsx +++ b/app/src/client/pages/AppsPage.tsx @@ -53,23 +53,22 @@ export function AppsPage({ } function statusFor(key: AppKey): { text: string; cls: string } { + if (!settings.apps.includes(key)) return { text: "off", cls: "status" }; if (desk.deskState?.authUrls?.[key]) { return { text: "one step left", cls: "status" }; } const err = desk.deskState?.appErrors?.[key]; if (err) return { text: `couldn't connect — ${err}`, cls: "status err" }; - const server = desk.mcp?.servers?.[key]; - if (!settings.apps.includes(key)) return { text: "off", cls: "status" }; if (key === "princetoncourses" && settings.apps.includes("junction")) { return { text: "covered by TigerJunction", cls: "status on" }; } - if (server?.state === "ready" || server?.state === "connected") { - return { text: "connected", cls: "status on" }; - } - if (server?.state === "failed") { - return { text: server.error ?? "connection failed", cls: "status err" }; + if (key === "gcal") { + return desk.deskState?.gcalReady + ? { text: "connected", cls: "status on" } + : { text: saving ? "connecting…" : "on", cls: "status" }; } - return { text: saving ? "connecting…" : "on", cls: "status" }; + // Engine connections open per turn, so "on" is the honest steady state. + return { text: "on", cls: "status on" }; } return ( diff --git a/app/src/client/pages/ChatPage.tsx b/app/src/client/pages/ChatPage.tsx index 9dd5155..f394495 100644 --- a/app/src/client/pages/ChatPage.tsx +++ b/app/src/client/pages/ChatPage.tsx @@ -81,12 +81,17 @@ export function ChatPage({ } }); - async function ensureSetup(force = false) { - if (!force && appliedRef.current === settingsHash) return; + /** + * Push settings to the agent. With `connect`, it also opens the MCP + * connections a turn needs — connections are held only while work runs + * (an idle Durable Object with live MCP clients never hibernates). + */ + async function ensureSetup(connect = false) { + if (!connect && appliedRef.current === settingsHash) return; appliedRef.current = settingsHash; try { await agent.ready; - await agent.call("setup", [settings]); + await agent.call("setup", [settings, { connect }]); } catch (err) { console.warn("PI setup failed", err); appliedRef.current = null; @@ -112,8 +117,7 @@ export function ChatPage({ title: firstTitle(messages) ?? text.slice(0, 48), at: Date.now(), }); - // Force a reconcile so connections made since the last message (like a - // Google Calendar consent finished on My apps) are live for this turn. + // Open this turn's MCP connections (released again when the turn ends). await ensureSetup(true); void sendMessage({ text }); } @@ -233,7 +237,7 @@ export function ChatPage({ } onRegenerate={ m.role === "assistant" && m.id === lastAssistantId - ? () => void regenerate() + ? () => void ensureSetup(true).then(() => regenerate()) : undefined } /> diff --git a/app/src/server/pi.ts b/app/src/server/pi.ts index b241b85..4eed388 100644 --- a/app/src/server/pi.ts +++ b/app/src/server/pi.ts @@ -22,6 +22,8 @@ export type PiState = { appErrors: Partial>; /** OAuth consent URLs awaiting the user (e.g. Google Calendar). */ authUrls?: Partial>; + /** Whether the desk holds Google tokens (chats can use the calendar). */ + gcalReady?: boolean; }; const CLAUDE_MODELS = new Set(["claude-opus-5", "claude-sonnet-5"]); @@ -35,7 +37,12 @@ const CAMPUS_MODEL = "@cf/zai-org/glm-4.7-flash"; * MCP endpoints, spoken over MCP v2 stateless streamable HTTP. */ export class Pi extends Think { - initialState: PiState = { settings: null, appErrors: {}, authUrls: {} }; + initialState: PiState = { + settings: null, + appErrors: {}, + authUrls: {}, + gcalReady: false, + }; /** Pure chat agent — no shell tool. Workspace file tools stay available. */ workspaceBash = false; @@ -101,7 +108,7 @@ export class Pi extends Think { * a reconnect of every app. */ @callable() - async setup(settings: PiSettings) { + async setup(settings: PiSettings, opts: { connect?: boolean } = {}) { // The Worker only routes a user to instances named `u--…`, so // requiring the same prefix here pins the MCP identity headers to the // signed-in user — settings.netid can't be spoofed sideways. @@ -148,6 +155,9 @@ export class Pi extends Think { const connectedIds = new Set(Object.keys(this.getMcpServers().servers)); for (const app of PI_APPS) { if (!enabled.has(app.key) || connectedIds.has(app.key)) continue; + // Engine connections are only opened right before a turn (see + // releaseIdleMcp for why); Google is handled here for the consent flow. + if (app.key !== "gcal" && !opts.connect) continue; try { if (app.key === "gcal") { if ( @@ -209,8 +219,54 @@ export class Pi extends Think { else appErrors.gcal = "couldn't start Google sign-in — try again"; } - this.setState({ settings, appErrors, authUrls }); - return { ok: true as const, appErrors, authUrls }; + const gcalReady = enabled.has("gcal") + ? this.isDesk() + ? await this.gcalTokensHas() + : await (await this.deskStub(settings.netid)).gcalTokensHas() + : false; + + // Nothing is about to run — don't sit on open connections. + if (!opts.connect) await this.releaseIdleMcp(); + + this.setState({ settings, appErrors, authUrls, gcalReady }); + return { ok: true as const, appErrors, authUrls, gcalReady }; + } + + /** + * A Durable Object holding live MCP connections never hibernates — it + * stays resident (and billed) around the clock even with zero traffic, + * while an identical object without them sleeps instantly. So PI keeps + * connections only for the duration of work: opened just before a turn or + * a direct tool call, released as soon as it ends, and swept on every + * wake in case a previous isolate left some behind. + */ + private async releaseIdleMcp(): Promise { + for (const [id, server] of Object.entries(this.getMcpServers().servers)) { + // Keep a Google connection that's mid-consent; its OAuth state lives + // on the connection row and the callback needs it. + if (id === "gcal" && server.state === "authenticating") continue; + try { + await this.removeMcpServer(id); + } catch (err) { + console.warn(`release ${id} failed`, err); + } + } + } + + override async onStart(props?: Record) { + await super.onStart(props); + await this.releaseIdleMcp(); + } + + override async onChatResponse(result: Parameters[0]) { + await super.onChatResponse(result); + await this.releaseIdleMcp(); + } + + override async onChatError(error: unknown, ctx?: Parameters[1]) { + const out = await super.onChatError(error, ctx); + await this.releaseIdleMcp(); + return out; } /** @@ -219,24 +275,33 @@ export class Pi extends Think { */ @callable() async callAppTool(app: AppKey, name: string, args: Record) { + const settings = this.getConfig(); + if (!settings?.apps.includes(app)) throw new Error(`${app} is switched off`); if (!this.getMcpServers().servers[app]) { - throw new Error(`${app} is not connected`); + await this.setup(settings, { connect: true }); + if (!this.getMcpServers().servers[app]) { + throw new Error(`${app} is not connected`); + } } - const result = (await this.mcp.callTool({ - serverId: app, - name, - arguments: args, - })) as { - isError?: boolean; - content?: Array<{ type: string; text?: string }>; - }; - const text = - result.content?.find((c) => c.type === "text" && c.text)?.text ?? ""; - if (result.isError) throw new Error(text || `${name} failed`); try { - return JSON.parse(text); - } catch { - return { text }; + const result = (await this.mcp.callTool({ + serverId: app, + name, + arguments: args, + })) as { + isError?: boolean; + content?: Array<{ type: string; text?: string }>; + }; + const text = + result.content?.find((c) => c.type === "text" && c.text)?.text ?? ""; + if (result.isError) throw new Error(text || `${name} failed`); + try { + return JSON.parse(text); + } catch { + return { text }; + } + } finally { + await this.releaseIdleMcp(); } } @@ -347,6 +412,35 @@ export class Pi extends Think { return url; } + /** TEMPORARY: what is keeping this object awake? */ + @callable() + async diag() { + const alarm = await this.ctx.storage.getAlarm(); + // eslint-disable-next-line @typescript-eslint/no-explicit-any + const self = this as any; + const sql = (q: string) => { + try { + return [...this.ctx.storage.sql.exec(q)]; + } catch (e) { + return [String(e)]; + } + }; + return { + now: Date.now(), + alarmAt: alarm, + alarmInMs: alarm == null ? null : alarm - Date.now(), + keepAliveRefs: self._keepAliveRefs, + pendingFiberRecovery: typeof self._hasPendingFiberRecovery === "function" ? self._hasPendingFiberRecovery() : "n/a", + schedules: sql("SELECT id, callback, type, time, running FROM cf_agents_schedules"), + facetRuns: sql("SELECT COUNT(*) AS n FROM cf_agents_facet_runs"), + tables: sql("SELECT name FROM sqlite_master WHERE type='table'").map((r: { name?: string }) => r.name), + mcp: Object.fromEntries( + Object.entries(this.getMcpServers().servers).map(([id, s]) => [id, s.state]) + ), + pendingMcp: Object.keys(self.mcp?._pendingConnections ?? {}), + }; + } + /** The per-user desk instance is the token authority for Google OAuth. */ private isDesk(): boolean { return this.name.endsWith("-desk");