diff --git a/README.md b/README.md index 1091a15..5fbdd71 100644 --- a/README.md +++ b/README.md @@ -56,7 +56,7 @@ This project grew from a small demo into an indispensable "productivity assistan | **Terminal Mode** | Smooth TUI in terminal | | **Smart Context Management** | Auto-compact long session context, configurable threshold, millisecond-level token estimation | | **Skills System** | Skills integration, auto-discover `SKILL.md`, activate by scenario | -| **Deep MCP Integration** | Local/remote MCP servers, OAuth login, runtime dynamic switching | +| **Deep MCP Integration** | Local/remote MCP servers, OAuth login, per-session dynamic switching | | **Enterprise-Grade Security** | Tool approval system (auto-approve/manual-approve), supports once/session/deny modes | | **OpenAI Compatible** | Works with any OpenAI-compatible API, flexible multi-Provider configuration | @@ -103,19 +103,18 @@ First run will guide you through Provider/Model setup and save config to `~/.mem ``` memo-code/ ├── packages/ -│ ├── core/ # Core logic: Session state machine, Config handling -│ ├── tools/ # Tool routing, MCP Client management, built-in tools (exec_command, read_text_file, apply_patch...) -│ ├── tui/ # Terminal runtime: CLI entry, interactive TUI -└── docs/ # Technical documentation +│ ├── core/ # Agent engine: session state machine, LLM/tool loop, built-in tools, MCP client, skills +│ └── tui/ # Terminal runtime: CLI entry, interactive TUI (Ink) +└── site/ # Documentation website (Next.js, static export) ``` **Technical Highlights:** -- **Architecture**: Clean Core / Tools / TUI separation, state-machine driven session management -- **Testing**: Core + Tools coverage > 70%, complete unit + integration tests +- **Architecture**: Core engine with integrated tool routing, thin TUI on top, state-machine driven session management +- **Testing**: Unit + integration tests, coverage threshold ≥70% - **Protocol**: Native MCP (Model Context Protocol) support, can integrate any MCP tool server - **Token Estimation**: Real-time context monitoring based on tiktoken, configurable auto-compaction strategy -- **Distribution**: npm package with hot-reloading without perception +- **Distribution**: Published to npm with version-driven auto releases via CI --- @@ -125,9 +124,11 @@ memo-code/ - `apply_patch` - Structured patch editing (`*** Begin Patch`/`*** End Patch`) - `read_text_file` / `read_media_file` / `read_files` / `write_file` / `edit_file` / `list_directory` / `search_files` - Filesystem read/write/search - `webfetch` - Paged web fetching with markdown extraction and policy guards -- MCP resource access - `list_mcp_resources`, `read_mcp_resource` +- MCP resource access - `list_mcp_resources`, `list_mcp_resource_templates`, `read_mcp_resource` - `update_plan` - Structured task progress management +- `read_skill` - Load skill instructions on demand - `get_memory` - Persistent memory reading +- Agent collaboration - `spawn_agent` / `send_message` / `followup_task` / `wait_agent` / `interrupt_agent` / `list_agents` --- @@ -148,8 +149,8 @@ base_url = "https://api.openai.com/v1" command = "npx" args = ["-y", "@modelcontextprotocol/server-github"] -# Skills -active_skills = ["./skills/doc-writing/SKILL.md"] +# Skills (absolute paths to SKILL.md files) +active_skills = ["/path/to/skills/doc-writing/SKILL.md"] ``` --- diff --git a/README.zh.md b/README.zh.md index 131565d..b64b2bd 100644 --- a/README.zh.md +++ b/README.zh.md @@ -52,7 +52,7 @@ Memo 诞生于一个简单的想法：**我想验证一个最简单的 Agent 是 | **终端模式** | 终端 TUI 交互流畅 | | **智能上下文管理** | 自动压缩长会话上下文，支持配置压缩阈值，毫秒级 token 估算 | | **Skills 技能系统** | Skills 技能集成，自动发现 `SKILL.md`，支持按场景激活 | -| **MCP 深度集成** | 支持本地/远程 MCP 服务器，OAuth 登录，运行时动态切换 | +| **MCP 深度集成** | 支持本地/远程 MCP 服务器，OAuth 登录，会话级动态切换 | | **企业级安全** | 工具分级审批机制（自动批准/手动批准），支持单次/会话/拒绝三种模式 | | **OpenAI 兼容** | 支持任意 OpenAI 兼容 API，灵活配置多 Provider 切换 | @@ -93,29 +93,30 @@ memo ``` memo-code/ ├── packages/ -│ ├── core/ # 核心逻辑：Session 状态机、Config 处理 -│ ├── tools/ # Tool 路由、MCP Client管理、内置工具实现（exec_command, read_text_file, apply_patch...） -│ ├── tui/ # 终端运行时：CLI 入口、交互式 TUI -└── docs/ # 技术文档 +│ ├── core/ # 核心引擎：Session 状态机、LLM/工具循环、内置工具、MCP 客户端、技能 +│ └── tui/ # 终端运行时：CLI 入口、交互式 TUI (Ink) +└── site/ # 文档网站（Next.js 静态导出） ``` **技术亮点：** -- **架构**：清晰的 Core / Tools / TUI 分层，状态机驱动会话管理 -- **测试**：Core + Tools 覆盖率 > 70%，完整的单元 + 集成测试 +- **架构**：核心引擎内置工具路由，TUI 薄壳，状态机驱动会话管理 +- **测试**：单元 + 集成测试，覆盖率门槛 ≥70% - **协议**：原生支持 MCP (Model Context Protocol)，可接入任意 MCP 工具服务器 - **Token 估算**：基于 tiktoken 的实时上下文监控，支持可配置的自动压缩策略 -- **分发**：npm 包热加载无感知 +- **分发**：发布至 npm，CI 版本驱动自动发版 ## 🔧 内置工具 - `exec_command` / `write_stdin` - 执行 Shell 命令 -- `apply_patch` - 字符串级代码编辑（单文件/批量） +- `apply_patch` - 结构化补丁编辑（`*** Begin Patch`/`*** End Patch`） - `read_text_file` / `read_media_file` / `read_files` / `write_file` / `edit_file` / `list_directory` / `search_files` - 文件系统读写与检索 - `webfetch` - 支持分页、Markdown 提取与策略防护的网页抓取 -- MCP 资源访问 - `list_mcp_resources`, `read_mcp_resource` +- MCP 资源访问 - `list_mcp_resources`, `list_mcp_resource_templates`, `read_mcp_resource` - `update_plan` - 结构化任务进度管理 +- `read_skill` - 按需加载技能指令 - `get_memory` - 持久化记忆读取 +- Agent 协作 - `spawn_agent` / `send_message` / `followup_task` / `wait_agent` / `interrupt_agent` / `list_agents` ## ⚙️ 配置示例 @@ -134,8 +135,8 @@ base_url = "https://api.openai.com/v1" command = "npx" args = ["-y", "@modelcontextprotocol/server-github"] -# Skills -active_skills = ["./skills/doc-writing/SKILL.md"] +# Skills（SKILL.md 的绝对路径） +active_skills = ["/path/to/skills/doc-writing/SKILL.md"] ``` --- diff --git a/package.json b/package.json index df224cc..356d13f 100644 --- a/package.json +++ b/package.json @@ -23,9 +23,9 @@ "start": "tsx packages/tui/src/cli.tsx", "build": "tsup", "dev": "tsup --watch", - "site:dev": "pnpm --filter @memo-code/site dev", - "site:build": "pnpm --filter @memo-code/site build", - "site:start": "pnpm --filter @memo-code/site start", + "site:dev": "pnpm --dir site dev", + "site:build": "pnpm --dir site build", + "site:start": "pnpm --dir site start", "format": "oxfmt packages site package.json pnpm-workspace.yaml tsconfig.json tsup.config.ts vitest.config.ts vitest.setup.ts .oxfmtrc.json", "format:check": "oxfmt --check packages site package.json pnpm-workspace.yaml tsconfig.json tsup.config.ts vitest.config.ts vitest.setup.ts .oxfmtrc.json", "lint": "oxlint packages site", @@ -67,7 +67,7 @@ "string-width": "^7.2.0", "toml": "^3.0.0", "turndown": "^7.2.2", - "undici": "^6.23.0", + "undici": "^6.28.0", "zod": "^4.3.6" } } diff --git a/packages/core/package.json b/packages/core/package.json index 9df47cf..4627948 100644 --- a/packages/core/package.json +++ b/packages/core/package.json @@ -26,7 +26,8 @@ "diff": "^8.0.3", "ignore": "^7.0.5", "js-tiktoken": "^1.0.21", - "minimatch": "^10.0.1", + "minimatch": "^10.2.6", + "undici": "^6.28.0", "zod": "^4.3.6" }, "devDependencies": { diff --git a/packages/core/src/agent/communication.test.ts b/packages/core/src/agent/communication.test.ts new file mode 100644 index 0000000..40c5738 --- /dev/null +++ b/packages/core/src/agent/communication.test.ts @@ -0,0 +1,19 @@ +import { describe, expect, test } from 'vitest' +import { InputQueue } from './communication' + +describe('InputQueue', () => { + test('wakes event-driven waiters without polling', async () => { + const queue = new InputQueue() + const waiting = queue.waitForActivity(1_000) + queue.enqueue({ author: '/root/a', recipient: '/root', content: 'done', triggerTurn: false }) + await expect(waiting).resolves.toBe('mailbox') + }) + + test('only drains a triggered batch when it contains trigger work', () => { + const queue = new InputQueue() + queue.enqueue({ author: '/root', recipient: '/root/a', content: 'note', triggerTurn: false }) + expect(queue.drainTriggeredBatch()).toEqual([]) + queue.enqueue({ author: '/root', recipient: '/root/a', content: 'continue', triggerTurn: true }) + expect(queue.drainTriggeredBatch().map((message) => message.content)).toEqual(['note', 'continue']) + }) +}) diff --git a/packages/core/src/agent/communication.ts b/packages/core/src/agent/communication.ts new file mode 100644 index 0000000..b544144 --- /dev/null +++ b/packages/core/src/agent/communication.ts @@ -0,0 +1,89 @@ +export type InterAgentCommunication = { + author: string + recipient: string + content: string + triggerTurn: boolean +} + +export type InputQueueActivity = 'mailbox' | 'timeout' | 'closed' | 'aborted' + +type ActivityWaiter = (activity: InputQueueActivity) => void + +export class InputQueue { + private messages: InterAgentCommunication[] = [] + private waiters = new Set() + private closed = false + + enqueue(communication: InterAgentCommunication): void { + if (this.closed) throw new Error('agent mailbox is closed') + this.messages.push(communication) + this.publish('mailbox') + } + + hasMessages(): boolean { + return this.messages.length > 0 + } + + hasTrigger(): boolean { + return this.messages.some((message) => message.triggerTurn) + } + + drainAll(): InterAgentCommunication[] { + if (this.messages.length === 0) return [] + const drained = this.messages + this.messages = [] + return drained + } + + drainTriggeredBatch(): InterAgentCommunication[] { + if (!this.hasTrigger()) return [] + return this.drainAll() + } + + async waitForActivity(timeoutMs: number, signal?: AbortSignal): Promise { + if (this.hasMessages()) return 'mailbox' + if (this.closed) return 'closed' + if (signal?.aborted) return 'aborted' + + return new Promise((resolve) => { + let settled = false + const finish = (activity: InputQueueActivity) => { + if (settled) return + settled = true + clearTimeout(timer) + signal?.removeEventListener('abort', onAbort) + this.waiters.delete(finish) + resolve(activity) + } + const onAbort = () => finish('aborted') + const timer = setTimeout(() => finish('timeout'), timeoutMs) + + this.waiters.add(finish) + signal?.addEventListener('abort', onAbort, { once: true }) + + // Recheck after subscribing so an enqueue cannot be missed at the boundary. + if (this.hasMessages()) finish('mailbox') + else if (this.closed) finish('closed') + }) + } + + close(): void { + if (this.closed) return + this.closed = true + this.publish('closed') + } + + private publish(activity: InputQueueActivity): void { + for (const waiter of this.waiters) waiter(activity) + } +} + +export function formatInterAgentCommunication(message: InterAgentCommunication): string { + return `\n${message.content}\n` +} + +function escapeAttribute(value: string): string { + return value.replace(/&/g, '&').replace(/"/g, '"').replace(//g, '>') +} diff --git a/packages/core/src/agent/compact_prompt.test.ts b/packages/core/src/agent/compact_prompt.test.ts index 7f9fafa..b7abf82 100644 --- a/packages/core/src/agent/compact_prompt.test.ts +++ b/packages/core/src/agent/compact_prompt.test.ts @@ -35,8 +35,9 @@ describe('compact_prompt', () => { const prompt = buildCompactionUserPrompt(messages) assert.ok(prompt.includes('[0] ASSISTANT (tool_calls: exec_command)')) assert.ok(prompt.includes('[1] TOOL (exec_command)')) - // Truncation keeps the tail (tool result/error at the end carries the info). - assert.ok(prompt.includes(`...${'x'.repeat(4_000)}`)) + // Long tool output is NOT truncated per-message: the budget selector is + // the only truncation point, so the compaction model sees full outputs. + assert.ok(prompt.includes(longToolOutput)) assert.ok(prompt.includes('Return only the summary body in plain text. Do not add markdown fences.')) }) @@ -93,13 +94,80 @@ describe('compact_prompt', () => { ) }) - test('selectCompactionMessages keeps the newest message even when it exceeds the budget', () => { + test('selectCompactionMessages truncates a newest message that alone exceeds the budget', () => { const messages: ChatMessage[] = [ { role: 'user', content: 'old' }, { role: 'assistant', content: 'x'.repeat(500) }, ] const selected = selectCompactionMessages(messages, 10, (text) => text.length) - assert.deepStrictEqual(selected, [messages[1]]) + assert.strictEqual(selected.length, 1) + const content = selected[0]?.content + assert.strictEqual(typeof content, 'string') + // Tail kept (with the truncation marker), head dropped. + assert.ok(String(content).startsWith('...')) + assert.ok(String(content).endsWith('x')) + assert.ok(String(content).length < 500) + }) + + test('plan from update_plan tool results survives into the compaction prompt', () => { + const planJson = JSON.stringify({ + message: 'Plan updated', + plan: [ + { step: 'Implement the parser', status: 'in_progress' }, + { step: 'Wire up the CLI flag', status: 'pending' }, + { step: 'Add tests for edge cases', status: 'pending' }, + ], + }) + const messages: ChatMessage[] = [ + { role: 'user', content: 'Refactor the parser' }, + { + role: 'assistant', + content: [{ type: 'text', text: 'Let me update the plan' }], + }, + { + role: 'tool', + content: [ + { + type: 'tool-result', + toolCallId: 'plan-1', + toolName: 'update_plan', + output: { type: 'text', value: planJson }, + }, + ], + }, + { role: 'user', content: 'Continue with step two' }, + ] + + const selected = selectCompactionMessages(messages, 10_000, (text) => text.length) + const prompt = buildCompactionUserPrompt(selected) + assert.ok(prompt.includes('Implement the parser'), 'plan steps must reach the compaction model') + assert.ok(prompt.includes('Wire up the CLI flag')) + assert.ok(prompt.includes('update_plan')) + }) + + test('selectCompactionMessages truncates oversized tool results keeping the tail', () => { + const messages: ChatMessage[] = [ + { + role: 'tool', + content: [ + { + type: 'tool-result', + toolCallId: 'call-1', + toolName: 'exec_command', + output: { type: 'text', value: `head-noise\n${'y'.repeat(1_000)}` }, + }, + ], + }, + ] + const selected = selectCompactionMessages(messages, 20, (text) => text.length) + assert.strictEqual(selected.length, 1) + const part = selected[0]?.content + assert.ok(Array.isArray(part)) + const value = part?.[0]?.type === 'tool-result' ? part[0].output.value : '' + assert.ok(String(value).startsWith('...')) + assert.ok(String(value).endsWith('y')) + assert.ok(String(value).length < 1_000) + assert.ok(!String(value).includes('head-noise'), 'tool output head is dropped') }) test('selectCompactionMessages returns empty for an empty array', () => { diff --git a/packages/core/src/agent/compact_prompt.ts b/packages/core/src/agent/compact_prompt.ts index e7555d2..2f34a9c 100644 --- a/packages/core/src/agent/compact_prompt.ts +++ b/packages/core/src/agent/compact_prompt.ts @@ -1,7 +1,5 @@ import type { ChatMessage } from '@memo/core/types' -const MAX_MESSAGE_CONTENT_CHARS = 4_000 - export const CONTEXT_COMPACTION_SYSTEM_PROMPT = `You are performing a CONTEXT CHECKPOINT COMPACTION. Create a handoff summary for another LLM that will resume the task. Include: @@ -16,13 +14,29 @@ export const CONTEXT_SUMMARY_PREFIX = 'Another language model started to solve this problem and produced a summary of its thinking process. You also have access to the state of the tools that were used by that language model. Use this to build on the work that has already been done and avoid duplicating work. Here is the summary produced by the other language model, use the information in this summary to assist with your own analysis:' function normalizeContent(content: string): string { - const compact = content.replace(/\r\n/g, '\n').trim() - if (compact.length <= MAX_MESSAGE_CONTENT_CHARS) { - return compact + return content.replace(/\r\n/g, '\n').trim() +} + +/** + * Truncate a single message to maxChars, keeping the tail (tool output carries + * its result/error at the end, while the head is usually echoes and noise). + * Only applied by selectCompactionMessages when one message alone exceeds the + * budget - regular transcripts are never truncated per-message. + */ +function truncateMessage(message: ChatMessage, maxChars: number): ChatMessage { + if (typeof message.content === 'string') { + if (message.content.length <= maxChars) return message + return { ...message, content: `...${message.content.slice(-maxChars)}` } + } + return { + ...message, + content: message.content.map((part) => { + if (part.type === 'tool-result' && part.output.type === 'text' && part.output.value.length > maxChars) { + return { ...part, output: { ...part.output, value: `...${part.output.value.slice(-maxChars)}` } } + } + return part + }), } - // Keep the tail: tool output carries its result/error at the end, while the - // head is usually command echoes and noise. - return `...${compact.slice(-MAX_MESSAGE_CONTENT_CHARS)}` } function messageToTranscriptLine(message: ChatMessage, index: number): string { @@ -60,9 +74,11 @@ export function isContextSummaryMessage(message: ChatMessage): boolean { /** * Drop the oldest messages so the serialized transcript fits within - * budgetTokens, keeping the newest message unconditionally. Returns the - * selected messages in their original order (indices are preserved — gaps - * mark the dropped messages). + * budgetTokens, keeping the newest message unconditionally. A single message + * that alone exceeds the budget is truncated to fit (tail kept), so the + * compaction request can never overflow the window. Returns the selected + * messages in their original order (indices are preserved — gaps mark the + * dropped messages). */ export function selectCompactionMessages( messages: ChatMessage[], @@ -80,10 +96,20 @@ export function selectCompactionMessages( if (!message) { continue } - const tokens = countTokens(messageToTranscriptLine(message, i)) + 1 // +1 for the '\n\n' separator + const line = messageToTranscriptLine(message, i) + const tokens = countTokens(line) + 1 // +1 for the '\n\n' separator if (selected.length > 0 && used + tokens > budgetTokens) { break } + if (selected.length === 0 && used + tokens > budgetTokens) { + // Newest message alone exceeds the budget: truncate it to fit. + // Linear scale from the estimated tokens gives a close-enough + // character budget for the tail. + const ratio = Math.min(1, Math.max(0, budgetTokens - used) / Math.max(1, tokens - 1)) + const maxChars = Math.max(1, Math.floor(line.length * ratio)) + selected.push(truncateMessage(message, maxChars)) + break + } selected.push(message) used += tokens if (used >= budgetTokens) { diff --git a/packages/core/src/agent/constants.ts b/packages/core/src/agent/constants.ts index acfde3b..ba26e79 100644 --- a/packages/core/src/agent/constants.ts +++ b/packages/core/src/agent/constants.ts @@ -3,7 +3,7 @@ import type { SessionMode } from '@memo/core/types' import type { ToolActionStatus } from '@memo/core/tools/approval' export const DEFAULT_SESSION_MODE: SessionMode = 'interactive' -export const DEFAULT_CONTEXT_WINDOW = 120_000 +export const DEFAULT_CONTEXT_WINDOW = 1_000_000 export const TOOL_ACTION_SUCCESS_STATUS: ToolActionStatus = 'success' export const TOOL_DISABLED_ERROR_MESSAGE = 'Tool usage is disabled in the current permission mode. Switch to /core/tools once or /core/tools full to enable tools.' diff --git a/packages/core/src/agent/control.ts b/packages/core/src/agent/control.ts new file mode 100644 index 0000000..4848cdf --- /dev/null +++ b/packages/core/src/agent/control.ts @@ -0,0 +1,302 @@ +import { randomUUID } from 'node:crypto' +import type { AgentActivity, AgentSession, HistoryEvent, HistorySink } from '@memo/core/types' +import { InputQueue, type InterAgentCommunication } from './communication' +import { buildForkHistory, buildSubagentSystemPrompt, parseForkTurns } from './fork' +import { AgentRegistry, ExecutionLimiter, ROOT_AGENT_PATH, agentPathDepth, type AgentMetadata } from './registry' +import { AgentRuntime } from './runtime' +import { deriveAgentStatusFromEvent, isFinalAgentStatus } from './status' + +export type CollabSessionBinding = { + agentId: string + agentPath: string + inputQueue: InputQueue + control: AgentControl +} + +export type CreateSubagentParams = { + sessionId: string + agentPath: string + initialHistory: AgentSession['history'] + systemPrompt: string + binding: CollabSessionBinding + statusSink: HistorySink +} + +export type AgentControlOptions = { + rootSessionId: string + rootSystemPrompt: string + rootHistorySinks: HistorySink[] + createSubagent: (params: CreateSubagentParams) => Promise + onAgentActivity?: (activity: AgentActivity) => void + maxRunningTurns?: number + maxDepth?: number +} + +function parsePositiveInteger(value: string | undefined, fallback: number): number { + const parsed = Number(value) + return Number.isInteger(parsed) && parsed > 0 ? parsed : fallback +} + +export class AgentControl { + private readonly registry = new AgentRegistry() + private readonly bindings = new Map() + private readonly sessions = new Map() + private readonly runtimes = new Map() + private readonly limiter: ExecutionLimiter + private readonly maxDepth: number + private shuttingDown = false + + constructor(private readonly options: AgentControlOptions) { + this.maxDepth = options.maxDepth ?? parsePositiveInteger(process.env.MEMO_SUBAGENT_MAX_DEPTH, 3) + const maxRunning = options.maxRunningTurns ?? parsePositiveInteger(process.env.MEMO_SUBAGENT_MAX_AGENTS, 4) + this.limiter = new ExecutionLimiter(maxRunning, () => { + queueMicrotask(() => this.wakePendingAgents()) + }) + this.registry.registerRoot(options.rootSessionId) + } + + createRootBinding(): CollabSessionBinding { + const existing = this.bindings.get(this.options.rootSessionId) + if (existing) return existing + const binding: CollabSessionBinding = { + agentId: this.options.rootSessionId, + agentPath: ROOT_AGENT_PATH, + inputQueue: new InputQueue(), + control: this, + } + this.bindings.set(binding.agentId, binding) + return binding + } + + attachRootSession(session: AgentSession): void { + this.sessions.set(this.options.rootSessionId, session) + } + + async spawnAgent( + sender: CollabSessionBinding, + params: { message: string; taskName: string; forkTurns?: string }, + ): Promise { + if (this.shuttingDown) throw new Error('agent tree is shutting down') + const parentSession = this.sessions.get(sender.agentId) + if (!parentSession) throw new Error(`parent agent is not loaded: ${sender.agentPath}`) + + const reservation = this.registry.reserve(sender.agentPath, params.taskName, this.maxDepth) + const permit = this.limiter.tryAcquire() + if (!permit) { + reservation.release() + throw new Error('subagent concurrency limit reached') + } + + const agentId = randomUUID() + const binding: CollabSessionBinding = { + agentId, + agentPath: reservation.agentPath, + inputQueue: new InputQueue(), + control: this, + } + const systemPrompt = buildSubagentSystemPrompt(this.options.rootSystemPrompt, binding.agentPath) + let session: AgentSession | undefined + try { + const initialHistory = buildForkHistory( + parentSession.history, + systemPrompt, + parseForkTurns(params.forkTurns), + ) + const statusSink: HistorySink = { + append: (event) => this.handleAgentEvent(agentId, event), + } + session = await this.options.createSubagent({ + sessionId: agentId, + agentPath: binding.agentPath, + initialHistory, + systemPrompt, + binding, + statusSink, + }) + const metadata: AgentMetadata = { + agentId, + agentPath: binding.agentPath, + taskName: params.taskName.trim(), + parentId: sender.agentId, + parentPath: sender.agentPath, + status: 'pending_init', + updatedAt: new Date().toISOString(), + } + this.registry.register(metadata, reservation) + this.bindings.set(agentId, binding) + this.sessions.set(agentId, session) + const runtime = new AgentRuntime(session, binding.inputQueue, this.limiter, (error) => { + void this.recordUnexpectedError(agentId, error) + }) + this.runtimes.set(agentId, runtime) + await this.publishActivity(metadata) + runtime.deliverWithPermit( + { + author: sender.agentPath, + recipient: binding.agentPath, + content: params.message, + triggerTurn: true, + }, + permit, + ) + return metadata + } catch (error) { + permit.release() + reservation.release() + if (session) await session.close().catch(() => {}) + this.runtimes.delete(agentId) + this.sessions.delete(agentId) + this.bindings.delete(agentId) + this.registry.remove(agentId) + throw error + } + } + + sendMessage(sender: CollabSessionBinding, target: string, content: string, triggerTurn: boolean): AgentMetadata { + const receiver = this.resolveTarget(target, sender.agentPath) + if (triggerTurn && receiver.agentPath === ROOT_AGENT_PATH) { + throw new Error('follow-up tasks cannot target the root agent') + } + const communication: InterAgentCommunication = { + author: sender.agentPath, + recipient: receiver.agentPath, + content, + triggerTurn, + } + const runtime = this.runtimes.get(receiver.agentId) + if (runtime) runtime.deliver(communication) + else this.bindings.get(receiver.agentId)?.inputQueue.enqueue(communication) + return receiver + } + + async waitForActivity(binding: CollabSessionBinding, timeoutMs: number, signal?: AbortSignal) { + return binding.inputQueue.waitForActivity(timeoutMs, signal) + } + + interruptAgent(sender: CollabSessionBinding, target: string): AgentMetadata { + const receiver = this.resolveTarget(target, sender.agentPath) + if (receiver.agentPath === ROOT_AGENT_PATH) throw new Error('root is not a spawned agent') + if (receiver.agentId === sender.agentId) throw new Error('an agent cannot interrupt itself') + this.runtimes.get(receiver.agentId)?.interrupt() + return receiver + } + + listAgents(sender: CollabSessionBinding, pathPrefix?: string): AgentMetadata[] { + let prefix = sender.agentPath + if (pathPrefix) { + prefix = pathPrefix.startsWith('/') ? pathPrefix : `${sender.agentPath}/${pathPrefix}` + } + return this.registry.list(prefix) + } + + async shutdownDescendants(rootAgentId: string): Promise { + if (rootAgentId !== this.options.rootSessionId || this.shuttingDown) return + this.shuttingDown = true + const descendants = this.registry + .list(ROOT_AGENT_PATH) + .sort((left, right) => agentPathDepth(right.agentPath) - agentPathDepth(left.agentPath)) + for (const agent of descendants) { + await this.runtimes.get(agent.agentId)?.shutdown() + this.runtimes.delete(agent.agentId) + this.sessions.delete(agent.agentId) + this.bindings.delete(agent.agentId) + this.registry.remove(agent.agentId) + } + this.bindings.get(rootAgentId)?.inputQueue.close() + } + + private resolveTarget(target: string, senderPath: string): AgentMetadata { + const receiver = this.registry.resolve(target.trim(), senderPath) + if (!receiver) throw new Error(`agent not found: ${target}`) + return receiver + } + + private wakePendingAgents(): void { + if (this.shuttingDown) return + for (const runtime of this.runtimes.values()) { + runtime.wakePending() + } + } + + private async handleAgentEvent(agentId: string, event: HistoryEvent): Promise { + const current = this.registry.getById(agentId) + if (!current) return + const nextStatus = deriveAgentStatusFromEvent(current, event) + const contextPercent = + event.type === 'context_usage' && typeof event.meta?.usage_percent === 'number' + ? event.meta.usage_percent + : undefined + const updated = this.registry.update(agentId, { + ...nextStatus, + ...(contextPercent === undefined ? {} : { contextPercent }), + }) + if (!updated) return + + if ( + event.type === 'turn_start' || + event.type === 'context_usage' || + event.type === 'turn_end' || + event.type === 'session_end' + ) { + await this.publishActivity(updated) + } + + if (event.type !== 'turn_end' || !isFinalAgentStatus(updated.status) || this.shuttingDown) return + this.notifyParentOfFinal(updated) + } + + private notifyParentOfFinal(updated: AgentMetadata): void { + const parent = updated.parentId ? this.bindings.get(updated.parentId) : undefined + if (!parent) return + const detail = updated.error ?? updated.lastMessage ?? '' + parent.inputQueue.enqueue({ + author: updated.agentPath, + recipient: parent.agentPath, + content: `\n${detail}\n`, + triggerTurn: false, + }) + } + + private async recordUnexpectedError(agentId: string, error: Error): Promise { + const updated = this.registry.update(agentId, { status: 'errored', error: error.message }) + if (!updated) return + await this.publishActivity(updated) + if (!this.shuttingDown) this.notifyParentOfFinal(updated) + } + + private async publishActivity(metadata: AgentMetadata): Promise { + const activity: AgentActivity = { + agentId: metadata.agentId, + agentPath: metadata.agentPath, + taskName: metadata.taskName, + parentId: metadata.parentId, + status: metadata.status, + contextPercent: metadata.contextPercent, + lastMessage: metadata.lastMessage, + error: metadata.error, + updatedAt: metadata.updatedAt, + } + try { + this.options.onAgentActivity?.(activity) + } catch { + // UI callbacks must not change agent lifecycle behavior. + } + const event: HistoryEvent = { + ts: new Date().toISOString(), + sessionId: this.options.rootSessionId, + type: 'agent_status', + content: activity.lastMessage, + meta: { + agent_id: activity.agentId, + agent_path: activity.agentPath, + task_name: activity.taskName, + parent_id: activity.parentId, + status: activity.status, + context_percent: activity.contextPercent, + error: activity.error, + updated_at: activity.updatedAt, + }, + } + await Promise.allSettled(this.options.rootHistorySinks.map((sink) => sink.append(event))) + } +} diff --git a/packages/core/src/agent/defaults.ts b/packages/core/src/agent/defaults.ts index c8ab5c1..2575416 100644 --- a/packages/core/src/agent/defaults.ts +++ b/packages/core/src/agent/defaults.ts @@ -35,15 +35,7 @@ export function filterMcpServersBySelection( return filtered } -/** - * Complete dependencies with default strategy (tools, callLLM, prompt, history sinks, tokenizer). - * Caller can provide only callbacks/overrides, rest use default implementations. - */ -export async function withDefaultDeps( - deps: AgentSessionDeps, - options: AgentSessionOptions, - sessionId: string, -): Promise<{ +export type ResolvedSessionDeps = { tools: ToolSet callLLM: CallLLM loadPrompt: () => Promise @@ -52,7 +44,18 @@ export async function withDefaultDeps( dispose: () => Promise historyFilePath?: string skillIndex: SkillIndex -}> { + createChildHistory: (sessionId: string) => { historySinks: HistorySink[]; historyFilePath?: string } +} + +/** + * Complete dependencies with default strategy (tools, callLLM, prompt, history sinks, tokenizer). + * Caller can provide only callbacks/overrides, rest use default implementations. + */ +export async function withDefaultDeps( + deps: AgentSessionDeps, + options: AgentSessionOptions, + sessionId: string, +): Promise { const loaded = await loadMemoConfig() const config = loaded.config @@ -146,5 +149,13 @@ export async function withDefaultDeps( tokenCounter: deps.tokenCounter ?? createTokenCounter(), historyFilePath: historyFilePath, skillIndex, + createChildHistory: (childSessionId) => { + if (deps.historySinks) return { historySinks: [] } + const childHistoryFilePath = buildSessionPath(sessionsDir, childSessionId) + return { + historySinks: [new JsonlHistorySink(childHistoryFilePath)], + historyFilePath: childHistoryFilePath, + } + }, } } diff --git a/packages/core/src/agent/fork.test.ts b/packages/core/src/agent/fork.test.ts new file mode 100644 index 0000000..866f98a --- /dev/null +++ b/packages/core/src/agent/fork.test.ts @@ -0,0 +1,28 @@ +import { describe, expect, test } from 'vitest' +import { buildForkHistory } from './fork' + +describe('buildForkHistory', () => { + test('keeps user turns and final assistant text while dropping tools and reasoning', () => { + const history = buildForkHistory( + [ + { role: 'system', content: 'parent' }, + { role: 'user', content: 'first' }, + { + role: 'assistant', + content: [ + { type: 'text', text: 'answer' }, + { type: 'reasoning', text: 'private' }, + ], + }, + { role: 'tool', content: [] }, + { role: 'user', content: 'second' }, + ], + 'child', + 1, + ) + expect(history).toEqual([ + { role: 'system', content: 'child' }, + { role: 'user', content: 'second' }, + ]) + }) +}) diff --git a/packages/core/src/agent/fork.ts b/packages/core/src/agent/fork.ts new file mode 100644 index 0000000..c105c12 --- /dev/null +++ b/packages/core/src/agent/fork.ts @@ -0,0 +1,48 @@ +import type { ChatMessage } from '@memo/core/types' + +export type ForkTurns = 'none' | 'all' | number + +export function parseForkTurns(value: string | undefined): ForkTurns { + const normalized = value?.trim().toLowerCase() || 'all' + if (normalized === 'none' || normalized === 'all') return normalized + const parsed = Number.parseInt(normalized, 10) + if (!/^\d+$/.test(normalized) || parsed <= 0) { + throw new Error('fork_turns must be `none`, `all`, or a positive integer string') + } + return parsed +} + +export function buildForkHistory( + parentHistory: ChatMessage[], + childSystemPrompt: string, + forkTurns: ForkTurns, +): ChatMessage[] { + if (forkTurns === 'none') return [{ role: 'system', content: childSystemPrompt }] + + let source = parentHistory + if (typeof forkTurns === 'number') { + const userPositions = parentHistory + .map((message, index) => (message.role === 'user' ? index : -1)) + .filter((index) => index >= 0) + const start = userPositions[Math.max(0, userPositions.length - forkTurns)] ?? parentHistory.length + source = parentHistory.slice(start) + } + + const inherited = source.flatMap((message): ChatMessage[] => { + if (message.role === 'user') return [message] + if (message.role !== 'assistant') return [] + if (typeof message.content === 'string') return message.content ? [message] : [] + + const text = message.content + .filter((part): part is Extract<(typeof message.content)[number], { type: 'text' }> => part.type === 'text') + .map((part) => part.text) + .join('') + return text ? [{ role: 'assistant', content: text }] : [] + }) + + return [{ role: 'system', content: childSystemPrompt }, ...inherited] +} + +export function buildSubagentSystemPrompt(parentSystemPrompt: string, agentPath: string): string { + return `${parentSystemPrompt}\n\nYou are sub-agent ${agentPath}. Work only on the assigned task. Preserve repository instructions, coordinate through the collaboration tools when needed, and return a concise result to your parent agent.` +} diff --git a/packages/core/src/agent/loop.ts b/packages/core/src/agent/loop.ts index 91e624d..ca11b09 100644 --- a/packages/core/src/agent/loop.ts +++ b/packages/core/src/agent/loop.ts @@ -48,6 +48,8 @@ import { import type { ToolApprovalHooks } from '@memo/core/tools/sdk_tools' import { runWithRuntimeContext } from '@memo/core/tools/runtime/context' import type { ToolExecutionContext } from '@memo/core/tools/sdk_tools' +import { formatInterAgentCommunication, type InterAgentCommunication } from '@memo/core/agent/communication' +import type { CollabSessionBinding } from '@memo/core/agent/control' import { createStepGate } from '@memo/core/tools/runtime/step_gate' import { isToolSkippedOutput, @@ -58,7 +60,7 @@ import { toToolHistoryMessage, } from './messages' -const DEFAULT_AUTO_COMPACT_THRESHOLD_PERCENT = 80 +const DEFAULT_AUTO_COMPACT_THRESHOLD_PERCENT = 90 const COMPACT_USER_MESSAGE_MAX_TOKENS = 20_000 export type SessionOperationKind = 'turn' | 'compact' @@ -84,6 +86,11 @@ type ActiveSessionOperation = { done: Promise } +export type AgentSessionRuntimeOptions = { + initialHistory?: ChatMessage[] + collab?: CollabSessionBinding +} + /** In-process conversation Session, implements multi-turn execution and log writing. */ export class AgentSessionImpl implements AgentSession { public title?: string @@ -111,6 +118,7 @@ export class AgentSessionImpl implements AgentSession { private toolPermissionMode: ToolPermissionMode | 'auto' = 'auto' /** Thinking override; undefined follows the provider model profile. */ private thinkingOverride: boolean | undefined + private collab?: CollabSessionBinding constructor( private deps: AgentSessionDeps & { @@ -121,10 +129,11 @@ export class AgentSessionImpl implements AgentSession { systemPrompt: string, tokenCounter: TokenCounter, historyFilePath?: string, + runtimeOptions: AgentSessionRuntimeOptions = {}, ) { this.id = options.sessionId || randomUUID() this.mode = options.mode || DEFAULT_SESSION_MODE - this.history = [{ role: 'system', content: systemPrompt }] + this.history = runtimeOptions.initialHistory?.slice() ?? [{ role: 'system', content: systemPrompt }] this.tokenCounter = tokenCounter this.sinks = deps.historySinks ?? [] this.hooks = buildHookRunners(deps) @@ -137,6 +146,7 @@ export class AgentSessionImpl implements AgentSession { mode: resolvedPermission.approvalMode, }) this.thinkingOverride = options.thinking + this.collab = runtimeOptions.collab } /** 运行时切换思考模式（undefined 恢复为跟随模型 profile）。 */ @@ -451,16 +461,17 @@ export class AgentSessionImpl implements AgentSession { private buildToolApprovalHooks(turn: number, step: number): ToolApprovalHooks { return { onApprovalRequest: async (request: ApprovalRequest) => { + const contextualRequest = this.contextualizeApprovalRequest(request) await runHook(this.hooks, 'onApprovalRequest', { sessionId: this.id, turn, step, - request, + request: contextualRequest, }) }, requestApproval: async (request: ApprovalRequest): Promise => { if (this.deps.requestApproval) { - return this.deps.requestApproval(request) + return this.deps.requestApproval(this.contextualizeApprovalRequest(request)) } return 'deny' }, @@ -476,8 +487,51 @@ export class AgentSessionImpl implements AgentSession { } } + private contextualizeApprovalRequest(request: ApprovalRequest): ApprovalRequest { + return { + ...request, + sessionId: this.id, + agentId: this.collab?.agentId, + agentPath: this.collab?.agentPath, + } + } + + private async recordAgentMessages( + turn: number, + messages: InterAgentCommunication[], + step?: number, + ): Promise { + if (messages.length === 0) return false + for (const message of messages) { + const content = formatInterAgentCommunication(message) + this.history.push({ role: 'user', content }) + await this.emitEvent('agent_message', { + turn, + step, + content, + role: 'user', + meta: { + author: message.author, + recipient: message.recipient, + trigger_turn: message.triggerTurn, + }, + }) + } + return true + } + + private async drainPendingAgentMessages(turn: number, step: number): Promise { + const messages = this.collab?.inputQueue.drainAll() ?? [] + await this.recordAgentMessages(turn, messages, step) + } + + private async drainTriggeredAgentMessages(turn: number, step: number): Promise { + const messages = this.collab?.inputQueue.drainTriggeredBatch() ?? [] + return this.recordAgentMessages(turn, messages, step) + } + private async maybeGenerateSessionTitle(turn: number, originalPrompt: string) { - if (turn !== 1 || this.title) return + if (turn !== 1 || this.title || (this.collab && this.collab.agentPath !== '/root')) return const title = fallbackSessionTitleFromPrompt(originalPrompt) this.title = title @@ -566,6 +620,7 @@ export class AgentSessionImpl implements AgentSession { // ReAct 主循环 for (let step = 0; ; step++) { + await this.drainPendingAgentMessages(turn, step) let estimatedPrompt = this.tokenCounter.countMessages(this.history) await this.emitContextUsage( turn, @@ -631,6 +686,7 @@ export class AgentSessionImpl implements AgentSession { toolsDisabled: this.toolsDisabled, gate: createStepGate(), skillIndex: this.deps.skillIndex, + collab: this.collab, } try { const llmResult = await this.deps.callLLM( @@ -639,13 +695,13 @@ export class AgentSessionImpl implements AgentSession { if (chunk) { receivedAssistantChunk = true } - this.deps.onAssistantStep?.(chunk, step) + this.deps.onAssistantStep?.(chunk, step, this.id) }, { signal: abortController.signal, toolContext, thinking: this.thinkingOverride, - onReasoningChunk: (chunk) => this.deps.onReasoningChunk?.(chunk, step), + onReasoningChunk: (chunk) => this.deps.onReasoningChunk?.(chunk, step, this.id), }, ) const normalized = normalizeLLMResponse(llmResult) @@ -702,7 +758,7 @@ export class AgentSessionImpl implements AgentSession { } if (!receivedAssistantChunk && assistantText) { - this.deps.onAssistantStep?.(assistantText, step) + this.deps.onAssistantStep?.(assistantText, step, this.id) } const textToolCall = @@ -998,6 +1054,9 @@ export class AgentSessionImpl implements AgentSession { // 无工具调用：文本即最终回复 if (toolUseBlocks.length === 0) { + if (await this.drainTriggeredAgentMessages(turn, step)) { + continue + } this.resetActionRepetition() const shouldFallbackFromPreviousText = toolUseBlocks.length === 0 && @@ -1068,6 +1127,7 @@ export class AgentSessionImpl implements AgentSession { stepCount: steps.length, durationMs: Date.now() - turnStartedAt, tokens: turnUsage, + error_message: errorMessage, protocol_violation_count: protocolViolationCount || undefined, }, }) @@ -1113,6 +1173,10 @@ export class AgentSessionImpl implements AgentSession { private async closeInternal() { await this.activeOperation?.done + if (this.collab?.agentPath === '/root') { + await this.collab.control.shutdownDescendants(this.collab.agentId) + } + this.collab?.inputQueue.close() // 空会话（从未 runTurn）不写 session_end，避免空会话落盘；sink 清理始终执行。 if (this.sessionStartEmitted) { await this.emitEvent('session_end', { @@ -1150,7 +1214,13 @@ export class AgentSessionImpl implements AgentSession { step: payload.step, content: payload.content, role: payload.role, - meta: payload.meta, + meta: this.collab + ? { + ...payload.meta, + agent_id: this.collab.agentId, + agent_path: this.collab.agentPath, + } + : payload.meta, }) await emitEventToSinks(event, this.sinks) } diff --git a/packages/core/src/agent/registry.test.ts b/packages/core/src/agent/registry.test.ts new file mode 100644 index 0000000..bd5813a --- /dev/null +++ b/packages/core/src/agent/registry.test.ts @@ -0,0 +1,24 @@ +import { describe, expect, test } from 'vitest' +import { AgentRegistry, ExecutionLimiter } from './registry' + +describe('AgentRegistry', () => { + test('reserves paths atomically and rolls failed spawns back', () => { + const registry = new AgentRegistry() + registry.registerRoot('root-id') + const reservation = registry.reserve('/root', 'review', 3) + expect(() => registry.reserve('/root', 'review', 3)).toThrow('already exists') + reservation.release() + expect(registry.reserve('/root', 'review', 3).agentPath).toBe('/root/review') + }) +}) + +describe('ExecutionLimiter', () => { + test('limits running turns independently of registered identities', () => { + const limiter = new ExecutionLimiter(1) + const first = limiter.tryAcquire() + expect(first).not.toBeNull() + expect(limiter.tryAcquire()).toBeNull() + first?.release() + expect(limiter.tryAcquire()).not.toBeNull() + }) +}) diff --git a/packages/core/src/agent/registry.ts b/packages/core/src/agent/registry.ts new file mode 100644 index 0000000..488ffb7 --- /dev/null +++ b/packages/core/src/agent/registry.ts @@ -0,0 +1,161 @@ +import type { AgentStatusSnapshot } from '@memo/core/agent/status' + +export const ROOT_AGENT_PATH = '/root' + +export type AgentMetadata = AgentStatusSnapshot & { + agentId: string + agentPath: string + taskName: string + parentId?: string + parentPath?: string + contextPercent?: number + updatedAt: string +} + +export function agentPathDepth(path: string): number { + if (path === ROOT_AGENT_PATH) return 0 + return path.split('/').filter(Boolean).length - 1 +} + +export function joinAgentPath(parentPath: string, taskName: string): string { + return `${parentPath}/${taskName}` +} + +export function validateTaskName(taskName: string): string { + const normalized = taskName.trim() + if (!/^[a-z0-9][a-z0-9_-]*$/i.test(normalized)) { + throw new Error('task_name must contain only letters, numbers, underscores, or hyphens') + } + return normalized +} + +export class SpawnReservation { + private active = true + + constructor( + readonly agentPath: string, + private releaseReservation: () => void, + ) {} + + commit(): void { + this.active = false + } + + release(): void { + if (!this.active) return + this.active = false + this.releaseReservation() + } +} + +export class AgentRegistry { + private byId = new Map() + private byPath = new Map() + private reservedPaths = new Set() + + registerRoot(agentId: string): void { + const now = new Date().toISOString() + const root: AgentMetadata = { + agentId, + agentPath: ROOT_AGENT_PATH, + taskName: 'root', + status: 'completed', + updatedAt: now, + } + this.byId.set(agentId, root) + this.byPath.set(ROOT_AGENT_PATH, agentId) + } + + reserve(parentPath: string, taskName: string, maxDepth: number): SpawnReservation { + const normalized = validateTaskName(taskName) + const path = joinAgentPath(parentPath, normalized) + if (agentPathDepth(path) > maxDepth) { + throw new Error(`subagent depth limit reached (${maxDepth})`) + } + if (this.byPath.has(path) || this.reservedPaths.has(path)) { + throw new Error(`agent path already exists: ${path}`) + } + this.reservedPaths.add(path) + return new SpawnReservation(path, () => this.reservedPaths.delete(path)) + } + + register(metadata: AgentMetadata, reservation: SpawnReservation): void { + if (reservation.agentPath !== metadata.agentPath || this.byPath.has(metadata.agentPath)) { + throw new Error(`agent path already exists: ${metadata.agentPath}`) + } + this.reservedPaths.delete(metadata.agentPath) + this.byId.set(metadata.agentId, metadata) + this.byPath.set(metadata.agentPath, metadata.agentId) + reservation.commit() + } + + getById(agentId: string): AgentMetadata | undefined { + return this.byId.get(agentId) + } + + getByPath(agentPath: string): AgentMetadata | undefined { + const id = this.byPath.get(agentPath) + return id ? this.byId.get(id) : undefined + } + + resolve(target: string, senderPath: string): AgentMetadata | undefined { + const direct = this.getById(target) + if (direct) return direct + if (target.startsWith('/')) return this.getByPath(target) + return this.getByPath(joinAgentPath(senderPath, target)) ?? this.getByPath(`${ROOT_AGENT_PATH}/${target}`) + } + + update( + agentId: string, + update: Partial & { contextPercent?: number }, + ): AgentMetadata | undefined { + const current = this.byId.get(agentId) + if (!current) return undefined + const next = { ...current, ...update, updatedAt: new Date().toISOString() } + this.byId.set(agentId, next) + return next + } + + list(pathPrefix = ROOT_AGENT_PATH): AgentMetadata[] { + return [...this.byId.values()] + .filter((agent) => agent.agentPath !== ROOT_AGENT_PATH) + .filter((agent) => agent.agentPath === pathPrefix || agent.agentPath.startsWith(`${pathPrefix}/`)) + .sort((left, right) => left.agentPath.localeCompare(right.agentPath)) + } + + remove(agentId: string): void { + const metadata = this.byId.get(agentId) + if (!metadata || metadata.agentPath === ROOT_AGENT_PATH) return + this.byId.delete(agentId) + this.byPath.delete(metadata.agentPath) + } +} + +export type ExecutionPermit = { release: () => void } + +export class ExecutionLimiter { + private active = 0 + + constructor( + private readonly capacity: number, + private readonly onRelease?: () => void, + ) {} + + tryAcquire(): ExecutionPermit | null { + if (this.active >= this.capacity) return null + this.active += 1 + let released = false + return { + release: () => { + if (released) return + released = true + this.active -= 1 + this.onRelease?.() + }, + } + } + + get activeCount(): number { + return this.active + } +} diff --git a/packages/core/src/agent/runtime.ts b/packages/core/src/agent/runtime.ts new file mode 100644 index 0000000..6e775b7 --- /dev/null +++ b/packages/core/src/agent/runtime.ts @@ -0,0 +1,87 @@ +import type { AgentSession } from '@memo/core/types' +import { formatInterAgentCommunication, type InterAgentCommunication, type InputQueue } from './communication' +import type { ExecutionLimiter, ExecutionPermit } from './registry' + +export class AgentRuntime { + private activeTurn: Promise | null = null + private shuttingDown = false + + constructor( + private readonly session: AgentSession, + readonly mailbox: InputQueue, + private readonly limiter: ExecutionLimiter, + private readonly onUnexpectedError: (error: Error) => void, + ) {} + + deliver(communication: InterAgentCommunication): void { + if (this.shuttingDown) throw new Error('agent is shutdown') + + if (communication.triggerTurn && !this.activeTurn) { + const permit = this.limiter.tryAcquire() + if (!permit) throw new Error('subagent concurrency limit reached') + this.mailbox.enqueue(communication) + this.startPendingTurn(permit) + return + } + + this.mailbox.enqueue(communication) + } + + deliverWithPermit(communication: InterAgentCommunication, permit: ExecutionPermit): void { + if (this.shuttingDown || this.activeTurn) { + permit.release() + throw new Error('agent cannot start its initial turn') + } + this.mailbox.enqueue(communication) + this.startPendingTurn(permit) + } + + wakePending(): boolean { + if (this.shuttingDown || this.activeTurn || !this.mailbox.hasTrigger()) return false + const permit = this.limiter.tryAcquire() + if (!permit) return false + this.startPendingTurn(permit) + return true + } + + interrupt(): void { + this.session.cancelCurrentTurn?.('interrupted by another agent') + } + + async shutdown(): Promise { + if (this.shuttingDown) { + await this.activeTurn + return + } + this.shuttingDown = true + this.mailbox.close() + this.session.cancelCurrentTurn?.('agent shutdown') + await this.activeTurn + await this.session.close() + } + + get isRunning(): boolean { + return this.activeTurn !== null + } + + private startPendingTurn(permit: ExecutionPermit): void { + const messages = this.mailbox.drainTriggeredBatch() + if (messages.length === 0) { + permit.release() + return + } + + const input = messages.map(formatInterAgentCommunication).join('\n\n') + const running = this.session + .runTurn(input) + .then(() => {}) + .catch((error: unknown) => { + this.onUnexpectedError(error instanceof Error ? error : new Error(String(error))) + }) + .finally(() => { + if (this.activeTurn === running) this.activeTurn = null + permit.release() + }) + this.activeTurn = running + } +} diff --git a/packages/core/src/agent/session.ts b/packages/core/src/agent/session.ts index 3f7974d..1604bae 100644 --- a/packages/core/src/agent/session.ts +++ b/packages/core/src/agent/session.ts @@ -3,6 +3,8 @@ import { withDefaultDeps } from '@memo/core/agent/defaults' import { DEFAULT_SESSION_MODE } from '@memo/core/agent/constants' import { AgentSessionImpl } from '@memo/core/agent/loop' import type { AgentSession, AgentSessionDeps, AgentSessionOptions } from '@memo/core/types' +import { AgentControl } from '@memo/core/agent/control' +import { createSubagentSession } from '@memo/core/agent/subagent' export { SessionBusyError, SessionClosedError, type SessionOperationKind } from '@memo/core/agent/loop' @@ -16,13 +18,37 @@ export async function createAgentSession( const sessionId = options.sessionId || randomUUID() const resolved = await withDefaultDeps(deps, { ...options, sessionId }, sessionId) const systemPrompt = await resolved.loadPrompt() + const collabEnabled = process.env.MEMO_ENABLE_COLLAB_TOOLS !== '0' + let control: AgentControl | undefined + if (collabEnabled) { + control = new AgentControl({ + rootSessionId: sessionId, + rootSystemPrompt: systemPrompt, + rootHistorySinks: resolved.historySinks, + onAgentActivity: deps.onAgentActivity, + createSubagent: (params) => + createSubagentSession({ + deps: resolved, + rootDeps: deps, + options: { ...options, mode: options.mode ?? DEFAULT_SESSION_MODE }, + sessionId: params.sessionId, + systemPrompt: params.systemPrompt, + initialHistory: params.initialHistory, + binding: params.binding, + statusSink: params.statusSink, + }), + }) + } + const collab = control?.createRootBinding() const session = new AgentSessionImpl( { ...(deps as AgentSessionDeps), ...resolved }, { ...options, sessionId, mode: options.mode ?? DEFAULT_SESSION_MODE }, systemPrompt, resolved.tokenCounter, resolved.historyFilePath, + { collab }, ) await session.init() + control?.attachRootSession(session) return session } diff --git a/packages/core/src/agent/session_hooks.test.ts b/packages/core/src/agent/session_hooks.test.ts index 00a9060..dc5442b 100644 --- a/packages/core/src/agent/session_hooks.test.ts +++ b/packages/core/src/agent/session_hooks.test.ts @@ -503,7 +503,6 @@ describe('session hooks & middleware', () => { assert.strictEqual(sessionStart?.meta?.cwd, process.cwd()) assert.strictEqual(sessionStart?.meta?.providerName, 'mock') assert.strictEqual(sessionStart?.meta?.modelName, 'mock-model') - assert.strictEqual(sessionStart?.meta?.contextWindow, 64_000) assert.strictEqual(sessionStart?.meta?.toolPermissionMode, 'once') assert.strictEqual(sessionStart?.meta?.thinking, true) assert.strictEqual(events.find((event) => event.type === 'turn_start')?.meta?.thinking, true) diff --git a/packages/core/src/agent/status.test.ts b/packages/core/src/agent/status.test.ts new file mode 100644 index 0000000..1d49387 --- /dev/null +++ b/packages/core/src/agent/status.test.ts @@ -0,0 +1,23 @@ +import { describe, expect, test } from 'vitest' +import { deriveAgentStatusFromEvent, isFinalAgentStatus, type AgentStatusSnapshot } from './status' + +describe('agent status', () => { + test('derives lifecycle status from session events', () => { + let status: AgentStatusSnapshot = { status: 'pending_init' } + status = deriveAgentStatusFromEvent(status, { type: 'turn_start' }) + expect(status.status).toBe('running') + status = deriveAgentStatusFromEvent(status, { type: 'final', content: 'done' }) + status = deriveAgentStatusFromEvent(status, { type: 'turn_end', meta: { status: 'ok' } }) + expect(status).toEqual({ status: 'completed', lastMessage: 'done' }) + expect(isFinalAgentStatus(status.status)).toBe(true) + }) + + test('keeps interrupted agents resumable', () => { + const status = deriveAgentStatusFromEvent( + { status: 'running' }, + { type: 'turn_end', meta: { status: 'cancelled' } }, + ) + expect(status.status).toBe('interrupted') + expect(isFinalAgentStatus(status.status)).toBe(false) + }) +}) diff --git a/packages/core/src/agent/status.ts b/packages/core/src/agent/status.ts new file mode 100644 index 0000000..0be956e --- /dev/null +++ b/packages/core/src/agent/status.ts @@ -0,0 +1,40 @@ +import type { HistoryEvent } from '@memo/core/types' + +export type AgentStatus = 'pending_init' | 'running' | 'interrupted' | 'completed' | 'errored' | 'shutdown' + +export type AgentStatusSnapshot = { + status: AgentStatus + lastMessage?: string + error?: string +} + +export function isFinalAgentStatus(status: AgentStatus): boolean { + return status === 'completed' || status === 'errored' || status === 'shutdown' +} + +export function deriveAgentStatusFromEvent( + current: AgentStatusSnapshot, + event: Pick, +): AgentStatusSnapshot { + if (event.type === 'final' && event.content !== undefined) { + return { ...current, lastMessage: event.content } + } + + if (event.type === 'turn_start') { + return { status: 'running', lastMessage: current.lastMessage } + } + + if (event.type === 'turn_end') { + const turnStatus = event.meta?.status + if (turnStatus === 'ok') return { status: 'completed', lastMessage: current.lastMessage } + if (turnStatus === 'cancelled') return { status: 'interrupted', lastMessage: current.lastMessage } + const error = typeof event.meta?.error_message === 'string' ? event.meta.error_message : undefined + return { status: 'errored', lastMessage: current.lastMessage, error } + } + + if (event.type === 'session_end') { + return { status: 'shutdown', lastMessage: current.lastMessage } + } + + return current +} diff --git a/packages/core/src/agent/subagent.ts b/packages/core/src/agent/subagent.ts new file mode 100644 index 0000000..a9bdebb --- /dev/null +++ b/packages/core/src/agent/subagent.ts @@ -0,0 +1,33 @@ +import type { ResolvedSessionDeps } from '@memo/core/agent/defaults' +import { AgentSessionImpl } from '@memo/core/agent/loop' +import type { CollabSessionBinding } from '@memo/core/agent/control' +import type { AgentSession, AgentSessionDeps, AgentSessionOptions, ChatMessage, HistorySink } from '@memo/core/types' + +export async function createSubagentSession(params: { + deps: ResolvedSessionDeps + rootDeps: AgentSessionDeps + options: AgentSessionOptions + sessionId: string + systemPrompt: string + initialHistory: ChatMessage[] + binding: CollabSessionBinding + statusSink: HistorySink +}): Promise { + const childHistory = params.deps.createChildHistory(params.sessionId) + const session = new AgentSessionImpl( + { + ...params.deps, + historySinks: [...childHistory.historySinks, params.statusSink], + requestApproval: params.rootDeps.requestApproval, + // Shared model/MCP tools are borrowed from the root; only the root owns their disposer. + dispose: undefined, + }, + { ...params.options, sessionId: params.sessionId }, + params.systemPrompt, + params.deps.tokenCounter, + childHistory.historyFilePath, + { initialHistory: params.initialHistory, collab: params.binding }, + ) + await session.init() + return session +} diff --git a/packages/core/src/api_types.ts b/packages/core/src/api_types.ts index aeb47a5..f83ec6d 100644 --- a/packages/core/src/api_types.ts +++ b/packages/core/src/api_types.ts @@ -87,8 +87,6 @@ export type SessionDetail = SessionListItem & { providerName?: string /** Model name recorded at session start, if any. */ modelName?: string - /** Context window (tokens) recorded at session start. */ - contextWindow?: number /** Tool permission mode recorded at session start. */ toolPermissionMode?: string /** Most recently recorded thinking override (undefined follows the model profile). */ diff --git a/packages/core/src/config/config.test.ts b/packages/core/src/config/config.test.ts index e62029b..c7fc39a 100644 --- a/packages/core/src/config/config.test.ts +++ b/packages/core/src/config/config.test.ts @@ -142,7 +142,7 @@ describe('mcp config serialization', () => { expect(text).toContain('active_skills = ["/tmp/.memo/skills/doc-writing/SKILL.md"]') expect(text).toContain('mcp_oauth_credentials_store_mode = "file"') expect(text).toContain('mcp_oauth_callback_port = 8765') - expect(text).toContain('auto_compact_threshold_percent = 80') + expect(text).toContain('auto_compact_threshold_percent = 90') expect(text).toContain('[model_profiles.gpt-5]') expect(text).toContain('supports_parallel_tool_calls = true') expect(text).toContain('context_window = 272000') @@ -287,7 +287,7 @@ model = "deepseek-chat" const loaded = await loadMemoConfig() expect(loaded.config.mcp_oauth_credentials_store_mode).toBe('auto') - expect(loaded.config.auto_compact_threshold_percent).toBe(80) + expect(loaded.config.auto_compact_threshold_percent).toBe(90) }) test('loadMemoConfig reads custom auto_compact_threshold_percent', async () => { @@ -325,7 +325,7 @@ model = "deepseek-chat" await writeFile(join(home, 'config.toml'), configText, 'utf-8') const loaded = await loadMemoConfig() - expect(loaded.config.auto_compact_threshold_percent).toBe(80) + expect(loaded.config.auto_compact_threshold_percent).toBe(90) }) test('loadMemoConfig ignores legacy max_prompt_tokens', async () => { @@ -408,7 +408,7 @@ describe('resolveContextWindowForProvider', () => { { name: 'openai', model: 'gpt-5', env_api_key: 'OPENAI_API_KEY' }, ) - expect(contextWindow).toBe(120000) + expect(contextWindow).toBe(1_000_000) }) test('matches keys case-insensitively and trims surrounding spaces', () => { diff --git a/packages/core/src/config/config.ts b/packages/core/src/config/config.ts index c99d1c5..b5eeecd 100644 --- a/packages/core/src/config/config.ts +++ b/packages/core/src/config/config.ts @@ -67,8 +67,8 @@ type ParsedMemoConfig = Omit, 'providers'> & { providers?: u const DEFAULT_MEMO_HOME = join(homedir(), '.memo') const DEFAULT_SESSIONS_DIR = 'sessions' -const DEFAULT_CONTEXT_WINDOW = 120000 -const DEFAULT_AUTO_COMPACT_THRESHOLD_PERCENT = 80 +const DEFAULT_CONTEXT_WINDOW = 1_000_000 +const DEFAULT_AUTO_COMPACT_THRESHOLD_PERCENT = 90 const DEFAULT_CONFIG: MemoConfig = { current_provider: 'deepseek', diff --git a/packages/core/src/features/history/history_parser.test.ts b/packages/core/src/features/history/history_parser.test.ts index 528e01c..645c0b7 100644 --- a/packages/core/src/features/history/history_parser.test.ts +++ b/packages/core/src/features/history/history_parser.test.ts @@ -12,7 +12,6 @@ function buildSampleLog(): string { cwd: '/tmp/demo', providerName: 'deepseek', modelName: 'deepseek-chat', - contextWindow: 64000, toolPermissionMode: 'once', thinking: true, }, @@ -81,7 +80,6 @@ describe('parseHistoryLogToSessionDetail', () => { const detail = parseHistoryLogToSessionDetail(buildSampleLog(), '/tmp/demo/s1.jsonl') assert.strictEqual(detail.providerName, 'deepseek') assert.strictEqual(detail.modelName, 'deepseek-chat') - assert.strictEqual(detail.contextWindow, 64000) assert.strictEqual(detail.toolPermissionMode, 'once') assert.strictEqual(detail.thinking, true) }) @@ -108,12 +106,11 @@ describe('parseHistoryLogToSessionDetail', () => { ts: '2026-02-15T10:00:00.000Z', sessionId: 'legacy', type: 'session_start', - meta: { cwd: '/tmp/demo', contextWindow: -1, thinking: 'yes' }, + meta: { cwd: '/tmp/demo', thinking: 'yes' }, }) const detail = parseHistoryLogToSessionDetail(log, '/tmp/demo/legacy.jsonl') assert.strictEqual(detail.providerName, undefined) - assert.strictEqual(detail.contextWindow, undefined) assert.strictEqual(detail.thinking, undefined) }) diff --git a/packages/core/src/features/history/history_parser.ts b/packages/core/src/features/history/history_parser.ts index 287bed9..370bb7d 100644 --- a/packages/core/src/features/history/history_parser.ts +++ b/packages/core/src/features/history/history_parser.ts @@ -22,7 +22,6 @@ type ParseResultState = { cwd: string providerName?: string modelName?: string - contextWindow?: number toolPermissionMode?: string thinking?: boolean startedAt: string @@ -279,7 +278,6 @@ export function parseHistoryLogToSessionDetail(raw: string, filePath: string): S cwd: '', providerName: undefined, modelName: undefined, - contextWindow: undefined, toolPermissionMode: undefined, thinking: undefined, startedAt: events[0]?.ts ?? fallbackNow, @@ -306,10 +304,6 @@ export function parseHistoryLogToSessionDetail(raw: string, filePath: string): S if (providerName) state.providerName = providerName const modelName = safeString(event.meta.modelName) if (modelName) state.modelName = modelName - const contextWindow = event.meta.contextWindow - if (typeof contextWindow === 'number' && Number.isFinite(contextWindow) && contextWindow > 0) { - state.contextWindow = contextWindow - } const toolPermissionMode = safeString(event.meta.toolPermissionMode) if (toolPermissionMode) state.toolPermissionMode = toolPermissionMode if (typeof event.meta.thinking === 'boolean') { @@ -503,7 +497,6 @@ export function parseHistoryLogToSessionDetail(raw: string, filePath: string): S events, providerName: state.providerName, modelName: state.modelName, - contextWindow: state.contextWindow, toolPermissionMode: state.toolPermissionMode, thinking: state.thinking, compactionSummary: state.compactionSummary, diff --git a/packages/core/src/features/history/history_sink.test.ts b/packages/core/src/features/history/history_sink.test.ts index 1f28b7f..a83b62f 100644 --- a/packages/core/src/features/history/history_sink.test.ts +++ b/packages/core/src/features/history/history_sink.test.ts @@ -1,5 +1,5 @@ import { describe, expect, test, beforeEach, afterEach } from 'vitest' -import { unlink, readFile } from 'node:fs/promises' +import { unlink, readFile, stat } from 'node:fs/promises' import { existsSync } from 'node:fs' import { join } from 'node:path' import { tmpdir } from 'node:os' @@ -40,6 +40,14 @@ describe('JsonlHistorySink', () => { expect(parsed.type === 'session_start').toBe(true) }) + test('writes session file with owner-only permissions', async () => { + const sink = new JsonlHistorySink(filePath) + await sink.append(createHistoryEvent({ sessionId: 'test-session', type: 'session_start' })) + + const mode = (await stat(filePath)).mode & 0o777 + expect(mode).toBe(0o600) + }) + test('creates parent directory if not exists', async () => { const nestedPath = join(tmpdir(), `memo-test-nested-${Date.now()}`, 'subdir', 'test.jsonl') const sink = new JsonlHistorySink(nestedPath) diff --git a/packages/core/src/features/history/history_sink.ts b/packages/core/src/features/history/history_sink.ts index 5c3db44..5e00181 100644 --- a/packages/core/src/features/history/history_sink.ts +++ b/packages/core/src/features/history/history_sink.ts @@ -24,7 +24,9 @@ export class JsonlHistorySink implements HistorySink { } this.writeQueue = this.writeQueue.then(async () => { await this.ensureDirectory() - await appendFile(this.filePath, `${JSON.stringify(event)}\n`, 'utf8') + // Session logs contain full conversations and tool output, which may + // include secrets; restrict to the owner like the OAuth credential file. + await appendFile(this.filePath, `${JSON.stringify(event)}\n`, { encoding: 'utf8', mode: 0o600 }) }) return this.writeQueue } diff --git a/packages/core/src/prompt/prompt.md b/packages/core/src/prompt/prompt.md index ebb3c1c..318b250 100644 --- a/packages/core/src/prompt/prompt.md +++ b/packages/core/src/prompt/prompt.md @@ -167,10 +167,11 @@ Use `get_memory` to retrieve persisted memory context for the current workflow: ## Subagent Collaboration -- Subagent tools (`spawn_agent`, `send_input`, `resume_agent`, `wait`, `close_agent`) do not require approval; treat their execution as dangerous and keep scope explicit +- Subagent tools (`spawn_agent`, `send_message`, `followup_task`, `wait_agent`, `interrupt_agent`, `list_agents`) do not require approval; keep delegated tasks narrow and explicit - Use subagents only for decomposable, well-scoped tasks; avoid recursive spawn loops -- Send concise task prompts, wait for completion, then summarize results back into the main thread -- Call `close_agent` for finished agents to release resources +- Use `send_message` to queue context without waking an idle agent; use `followup_task` when it should continue immediately +- `wait_agent` waits for mailbox activity; the result is injected into the next model request instead of returned by the tool +- Use `interrupt_agent` to stop a current turn without destroying the agent's conversation history ## Tool Call Discipline (CRITICAL) diff --git a/packages/core/src/tools/approval/constants.ts b/packages/core/src/tools/approval/constants.ts index 6a1ff54..a0bca98 100644 --- a/packages/core/src/tools/approval/constants.ts +++ b/packages/core/src/tools/approval/constants.ts @@ -16,11 +16,12 @@ export const DEFAULT_TOOL_RISK_LEVELS: Record = { read_files: 'read', list_directory: 'read', search_files: 'read', - wait: 'read', + wait_agent: 'read', spawn_agent: 'read', - send_input: 'read', - resume_agent: 'read', - close_agent: 'read', + send_message: 'read', + followup_task: 'read', + interrupt_agent: 'read', + list_agents: 'read', // Write tools: require approval. apply_patch: 'write', @@ -40,10 +41,11 @@ export const DEFAULT_TOOL_RISK_LEVELS: Record = { */ export const ALWAYS_AUTO_APPROVE_TOOLS = new Set([ 'spawn_agent', - 'send_input', - 'resume_agent', - 'wait', - 'close_agent', + 'send_message', + 'followup_task', + 'wait_agent', + 'interrupt_agent', + 'list_agents', ]) /** Risk level ordering for comparisons. */ diff --git a/packages/core/src/tools/approval/types.ts b/packages/core/src/tools/approval/types.ts index f004bcf..3cb2b6d 100644 --- a/packages/core/src/tools/approval/types.ts +++ b/packages/core/src/tools/approval/types.ts @@ -31,6 +31,9 @@ export interface ApprovalRequest { fingerprint: ApprovalKey riskLevel: RiskLevel reason: string + sessionId?: string + agentId?: string + agentPath?: string } /** Approval manager configuration */ diff --git a/packages/core/src/tools/index.test.ts b/packages/core/src/tools/index.test.ts index f029061..379a49d 100644 --- a/packages/core/src/tools/index.test.ts +++ b/packages/core/src/tools/index.test.ts @@ -12,10 +12,11 @@ describe('toolkit defaults', () => { vi.resetModules() const mod = await import('./index') assert.ok(mod.TOOLKIT.spawn_agent) - assert.ok(mod.TOOLKIT.send_input) - assert.ok(mod.TOOLKIT.resume_agent) - assert.ok(mod.TOOLKIT.wait) - assert.ok(mod.TOOLKIT.close_agent) + assert.ok(mod.TOOLKIT.send_message) + assert.ok(mod.TOOLKIT.followup_task) + assert.ok(mod.TOOLKIT.wait_agent) + assert.ok(mod.TOOLKIT.interrupt_agent) + assert.ok(mod.TOOLKIT.list_agents) }) test('allows explicit disabling collab tools via MEMO_ENABLE_COLLAB_TOOLS=0', async () => { @@ -23,9 +24,10 @@ describe('toolkit defaults', () => { vi.resetModules() const mod = await import('./index') assert.strictEqual(mod.TOOLKIT.spawn_agent, undefined) - assert.strictEqual(mod.TOOLKIT.send_input, undefined) - assert.strictEqual(mod.TOOLKIT.resume_agent, undefined) - assert.strictEqual(mod.TOOLKIT.wait, undefined) - assert.strictEqual(mod.TOOLKIT.close_agent, undefined) + assert.strictEqual(mod.TOOLKIT.send_message, undefined) + assert.strictEqual(mod.TOOLKIT.followup_task, undefined) + assert.strictEqual(mod.TOOLKIT.wait_agent, undefined) + assert.strictEqual(mod.TOOLKIT.interrupt_agent, undefined) + assert.strictEqual(mod.TOOLKIT.list_agents, undefined) }) }) diff --git a/packages/core/src/tools/index.ts b/packages/core/src/tools/index.ts index d733926..8246a54 100644 --- a/packages/core/src/tools/index.ts +++ b/packages/core/src/tools/index.ts @@ -20,7 +20,14 @@ import { updatePlanTool } from '@memo/core/tools/tools/update_plan' import { getMemoryTool } from '@memo/core/tools/tools/get_memory' import { readSkillTool } from '@memo/core/tools/tools/read_skill' import { webfetchTool } from '@memo/core/tools/tools/webfetch' -import { closeAgentTool, resumeAgentTool, sendInputTool, spawnAgentTool, waitTool } from '@memo/core/tools/tools/collab' +import { + followupTaskTool, + interruptAgentTool, + listAgentsTool, + sendMessageTool, + spawnAgentTool, + waitAgentTool, +} from '@memo/core/tools/tools/collab' function buildCodexTools(): Record { const tools: Record = {} @@ -62,10 +69,11 @@ function buildCodexTools(): Record { if (collabEnabled) { tools.spawn_agent = spawnAgentTool - tools.send_input = sendInputTool - tools.resume_agent = resumeAgentTool - tools.wait = waitTool - tools.close_agent = closeAgentTool + tools.send_message = sendMessageTool + tools.followup_task = followupTaskTool + tools.wait_agent = waitAgentTool + tools.interrupt_agent = interruptAgentTool + tools.list_agents = listAgentsTool } return tools diff --git a/packages/core/src/tools/router/mcp/oauth.runtime.test.ts b/packages/core/src/tools/router/mcp/oauth.runtime.test.ts index 33d9454..9e8bca2 100644 --- a/packages/core/src/tools/router/mcp/oauth.runtime.test.ts +++ b/packages/core/src/tools/router/mcp/oauth.runtime.test.ts @@ -121,6 +121,30 @@ describe('oauth runtime helpers', () => { await expect(openExternalUrl('https://example.com/oauth/authorize')).rejects.toThrow('spawn failed') }) + test('openExternalUrl on win32 uses PowerShell Start-Process', async () => { + const child = spawnChild() + spawnMock.mockReturnValue(child) + const platformSpy = vi.spyOn(process, 'platform', 'get').mockReturnValue('win32') + + await openExternalUrl('https://example.com/oauth/authorize') + + expect(spawnMock).toHaveBeenCalledWith( + 'powershell.exe', + ['-NoProfile', '-NonInteractive', '-Command', "Start-Process 'https://example.com/oauth/authorize'"], + expect.objectContaining({ detached: true }), + ) + platformSpy.mockRestore() + }) + + test('openExternalUrl on win32 rejects URLs containing quotes', async () => { + const platformSpy = vi.spyOn(process, 'platform', 'get').mockReturnValue('win32') + + await expect(openExternalUrl('https://example.com/a"b')).rejects.toThrow('quote') + + expect(spawnMock).not.toHaveBeenCalled() + platformSpy.mockRestore() + }) + test('createRuntimeMcpOAuthProvider returns null when no stored token exists', async () => { await withTempMemoHome(async (home) => { const provider = await createRuntimeMcpOAuthProvider({ @@ -247,6 +271,43 @@ describe('oauth login flow', () => { ).rejects.toThrow('does not advertise OAuth support') }) + test('callback error page escapes attacker-controlled error text', async () => { + await withTempMemoHome(async (home) => { + installDiscoveryFetch() + spawnMock.mockReturnValue(spawnChild()) + + let redirectUrl = '' + authMock.mockImplementationOnce(async (provider: any) => { + redirectUrl = provider.redirectUrl + return 'REDIRECT' + }) + authMock.mockResolvedValueOnce('REDIRECT') + + const loginRejection = loginMcpServerOAuth({ + serverName: 'remote', + config: { type: 'streamable_http', url: TEST_URL }, + timeoutMs: 5_000, + settings: { memoHome: home, storeMode: 'file' }, + }).then( + () => null, + (error: unknown) => error, + ) + + await vi.waitFor(() => { + assert.ok(redirectUrl, 'callback server did not come up') + }) + + const res = await fetch(`${redirectUrl}?error=bad&error_description=`) + const body = await res.text() + + assert.strictEqual(res.status, 400) + expect(body).toContain('<script>') + expect(body).not.toContain('