diff --git a/kun/README.md b/kun/README.md index bbe80ce53..45358af1d 100644 --- a/kun/README.md +++ b/kun/README.md @@ -220,6 +220,13 @@ Shape: "transport": "streamable-http", "url": "https://mcp.example.com/mcp", "headers": { "authorization": "Bearer " }, + "oauth": { + "enabled": true, + "clientName": "Kun", + "clientId": "", + "clientSecret": "", + "scopes": ["docs.readonly"] + }, "trustScope": "user", "timeoutMs": 30000 } @@ -279,7 +286,7 @@ See `../docs/KUN_CONFIG.md` for the detailed file layout and examples. Feature flags are intentionally explicit: -- `capabilities.mcp` starts configured MCP clients and imports their tools into the dynamic registry. Workspace-scoped servers require `trustedWorkspaceRoots`. +- `capabilities.mcp` starts configured MCP clients and imports their tools into the dynamic registry. Workspace-scoped servers require `trustedWorkspaceRoots`. Remote HTTP/SSE servers can use `oauth`; Kun stores OAuth tokens under the data directory instead of the config file. Use `GET /v1/mcp/oauth` to inspect redacted OAuth state and `DELETE /v1/mcp/oauth/{serverId}` to clear a server's saved authorization. - `serve.mcpSearch` can collapse a large MCP catalog into four entry points: `mcp_search`, `mcp_describe`, `mcp_call`, and `mcp_refresh_catalog`. When the catalog is too large, the model searches for relevant tools first, then describes and calls the exact tool instead of carrying every MCP schema on every turn. - `serve.tokenEconomy` / `tokenEconomyMode` compresses tool descriptions, tool results, and history context while preserving code, paths, commands, URLs, errors, and other high-value signals. - `contextCompaction` controls fallback long-thread compaction thresholds and summary behavior. Per-model thresholds live in `models.profiles`. Compaction preserves goals, constraints, decisions, touched files, tool outcomes, and unresolved next steps. diff --git a/kun/README.zh-CN.md b/kun/README.zh-CN.md index 7458fd6a3..86f02d57f 100644 --- a/kun/README.zh-CN.md +++ b/kun/README.zh-CN.md @@ -188,6 +188,13 @@ Kun 使用 JSON 配置文件管理运行时行为，避免重建后重配或硬 "transport": "streamable-http", "url": "https://mcp.example.com/mcp", "headers": { "authorization": "Bearer " }, + "oauth": { + "enabled": true, + "clientName": "Kun", + "clientId": "", + "clientSecret": "", + "scopes": ["docs.readonly"] + }, "trustScope": "user", "timeoutMs": 30000 } @@ -235,7 +242,7 @@ Kun 默认使用混合存储：`threads/{threadId}/messages.jsonl` 与 `events.j 功能开关是显式设计： -- `capabilities.mcp` 启动配置化 MCP 客户端并将工具加入动态注册表；工作区级服务器要求设置 `trustedWorkspaceRoots`。 +- `capabilities.mcp` 启动配置化 MCP 客户端并将工具加入动态注册表；工作区级服务器要求设置 `trustedWorkspaceRoots`。远程 HTTP/SSE MCP 可配置 `oauth`，Kun 会把 OAuth token 存在数据目录下，而不是写进 config。使用 `GET /v1/mcp/oauth` 可查看脱敏后的 OAuth 状态，使用 `DELETE /v1/mcp/oauth/{serverId}` 可清除某个服务保存的授权。 - `serve.mcpSearch` 可把大量 MCP 工具收敛为 `mcp_search`、`mcp_describe`、`mcp_call` 和 `mcp_refresh_catalog` 四个入口；当工具目录过大时，模型先检索意图相关工具，再描述和调用具体工具，避免每轮都携带完整 MCP schema。 - `serve.tokenEconomy` / `tokenEconomyMode` 会压缩工具描述、工具结果和历史上下文；保留代码、路径、命令、URL、错误信号等高价值信息，同时省掉重复、超长或二进制 payload。 - `contextCompaction` 控制长会话压缩的兜底阈值和摘要方式；模型级阈值写在 `models.profiles`。压缩时保留目标、约束、决策、已触碰文件、工具结果和未解决事项。 diff --git a/kun/src/adapters/model/compat-model-client.ts b/kun/src/adapters/model/compat-model-client.ts index 87b86e6f7..82f087429 100644 --- a/kun/src/adapters/model/compat-model-client.ts +++ b/kun/src/adapters/model/compat-model-client.ts @@ -19,6 +19,7 @@ import { type ModelEndpointFormat } from '../../contracts/model-endpoint-format.js' import { createProxyFetch } from './proxy-fetch.js' +import { wrapUntrustedContent } from '../../security/untrusted-content.js' /** * Configuration for the compatible HTTP model client. Chat @@ -695,6 +696,9 @@ export class CompatModelClient implements ModelClient { if (request.attachmentTextFallbacks?.length) { attachTextFallbacksToLatestUserMessage(out, request.attachmentTextFallbacks) } + if (request.attachmentDocuments?.length) { + attachDocumentsToLatestUserMessage(out, request.attachmentDocuments) + } return normalizeThinkingAssistantMessages(healToolMessagePairs(out), thinkingMode) } @@ -2649,6 +2653,46 @@ function attachTextFallbacksToLatestUserMessage( } } +function attachDocumentsToLatestUserMessage( + messages: ChatMessage[], + documents: NonNullable +): void { + const text = documents.map(formatAttachmentDocument).join('\n\n') + for (let index = messages.length - 1; index >= 0; index -= 1) { + const message = messages[index] + if (message.role !== 'user') continue + if (typeof message.content === 'string') { + message.content = message.content ? `${message.content}\n\n${text}` : text + return + } + if (Array.isArray(message.content)) { + message.content.push({ type: 'text', text }) + return + } + message.content = text + return + } +} + +function formatAttachmentDocument( + document: NonNullable[number] +): string { + return [ + '[Attached document]', + `Name: ${document.name}`, + `FilePath: ${document.localFilePath ?? 'unknown'}`, + `MIME: ${document.mimeType}`, + ...(document.pageCount ? [`Pages: ${document.pageCount}`] : []), + ...(document.truncated ? ['Note: text truncated to fit the context limit'] : []), + 'Content:', + wrapUntrustedContent({ + content: document.text, + source: { kind: 'document', label: document.name } + }), + '[/Attached document]' + ].join('\n') +} + function formatAttachmentTextFallback( attachment: NonNullable[number] ): string { diff --git a/kun/src/adapters/tool/artifact-tool.test.ts b/kun/src/adapters/tool/artifact-tool.test.ts new file mode 100644 index 000000000..fa46a53af --- /dev/null +++ b/kun/src/adapters/tool/artifact-tool.test.ts @@ -0,0 +1,54 @@ +import { describe, expect, it, vi } from 'vitest' +import { createReadArtifactTool } from './artifact-tool.js' +import { InMemoryArtifactStore } from '../../artifacts/artifact-store.js' +import type { ToolHostContext } from '../../ports/tool-host.js' + +function context(artifactStore?: ToolHostContext['artifactStore']): ToolHostContext { + return { + threadId: 't', turnId: 'tn', workspace: '/ws', approvalPolicy: 'auto', + abortSignal: new AbortController().signal, awaitApproval: vi.fn(async () => 'allow' as const), + ...(artifactStore ? { artifactStore } : {}) + } +} + +describe('read_artifact tool', () => { + it('reads full content + metadata by id', async () => { + const store = new InMemoryArtifactStore(() => 't0') + const { meta } = await store.put({ content: 'hello world', source: 'mcp', origin: 'docs' }) + const tool = createReadArtifactTool() + const result = await tool.execute({ artifactId: meta.id }, context(store)) + expect(result.output).toMatchObject({ content: 'hello world', source: 'mcp', origin: 'docs' }) + }) + + it('reads a line range', async () => { + const store = new InMemoryArtifactStore() + const { meta } = await store.put({ content: 'l1\nl2\nl3\nl4' }) + const tool = createReadArtifactTool() + const result = await tool.execute({ artifactId: meta.id, startLine: 2, endLine: 3 }, context(store)) + expect(result.output).toMatchObject({ content: 'l2\nl3', range: { startLine: 2, endLine: 3 } }) + }) + + it('bounds a no-range read and returns a cursor for a large artifact', async () => { + const { ARTIFACT_MAX_READ_BYTES } = await import('../../artifacts/artifact-store.js') + const store = new InMemoryArtifactStore(() => 't0') + const { meta } = await store.put({ content: 'y'.repeat(ARTIFACT_MAX_READ_BYTES + 1_000) }) + const tool = createReadArtifactTool() + const result = await tool.execute({ artifactId: meta.id }, context(store)) + const out = result.output as Record + expect(Buffer.byteLength(String(out.content), 'utf8')).toBe(ARTIFACT_MAX_READ_BYTES) + expect(out.truncated).toBe(true) + expect(out.nextOffset).toBe(ARTIFACT_MAX_READ_BYTES) + }) + + it('errors for an unknown artifact', async () => { + const tool = createReadArtifactTool() + const result = await tool.execute({ artifactId: 'art_missing' }, context(new InMemoryArtifactStore())) + expect(result.isError).toBe(true) + }) + + it('errors when the artifact store is unavailable', async () => { + const tool = createReadArtifactTool() + const result = await tool.execute({ artifactId: 'art_x' }, context()) + expect(result.isError).toBe(true) + }) +}) diff --git a/kun/src/adapters/tool/artifact-tool.ts b/kun/src/adapters/tool/artifact-tool.ts new file mode 100644 index 000000000..9dacf77c6 --- /dev/null +++ b/kun/src/adapters/tool/artifact-tool.ts @@ -0,0 +1,69 @@ +/** + * Agent-callable artifact reader (P0 #5). + * + * Large tool results are offloaded to the content-addressed artifact store and + * the model only sees a bounded summary + an artifact id. This tool lets the + * model pull the rest on demand — the full content, a byte range, or a line + * range — instead of forcing every big payload into context. Read-only. + */ + +import { LocalToolHost, type LocalTool } from './local-tool-host.js' +import { readArtifactBounded } from '../../artifacts/artifact-store.js' + +export function createReadArtifactTool(): LocalTool { + return LocalToolHost.defineTool({ + name: 'read_artifact', + description: + 'Read a stored artifact (large tool output) by id. Optionally fetch only a slice: ' + + 'startLine/endLine for a 1-indexed inclusive line range, or offset/length for a UTF-8 byte range. ' + + 'Use the artifact id returned by a previous tool result (e.g. stdoutArtifactId).', + inputSchema: { + type: 'object', + properties: { + artifactId: { type: 'string' }, + offset: { type: 'number' }, + length: { type: 'number' }, + startLine: { type: 'number' }, + endLine: { type: 'number' } + }, + required: ['artifactId'], + additionalProperties: false + }, + policy: 'auto', + execute: async (args, context) => { + if (!context.artifactStore) { + return { output: { error: 'artifact store is not available in this runtime' }, isError: true } + } + const artifactId = typeof args.artifactId === 'string' ? args.artifactId.trim() : '' + if (!artifactId) return { output: { error: 'artifactId is required' }, isError: true } + const meta = await context.artifactStore.stat(artifactId) + if (!meta) return { output: { error: `artifact not found: ${artifactId}` }, isError: true } + const range = { + ...(typeof args.offset === 'number' ? { offset: args.offset } : {}), + ...(typeof args.length === 'number' ? { length: args.length } : {}), + ...(typeof args.startLine === 'number' ? { startLine: args.startLine } : {}), + ...(typeof args.endLine === 'number' ? { endLine: args.endLine } : {}) + } + // Always read through the bounded reader: a request with no range, or a + // range larger than the cap, is clamped to <=1 MiB / <=2000 lines and a + // cursor is returned so the model pages rather than pulling a huge result + // into context. + const bounded = await readArtifactBounded(context.artifactStore, artifactId, meta, range) + if (bounded === null) return { output: { error: `artifact content not found: ${artifactId}` }, isError: true } + return { + output: { + artifactId, + byteSize: meta.byteSize, + lineCount: meta.lineCount, + ...(meta.source ? { source: meta.source } : {}), + ...(meta.origin ? { origin: meta.origin } : {}), + range: bounded.range, + truncated: bounded.truncated, + ...(bounded.nextOffset !== undefined ? { nextOffset: bounded.nextOffset } : {}), + ...(bounded.nextStartLine !== undefined ? { nextStartLine: bounded.nextStartLine } : {}), + content: bounded.content + } + } + } + }) +} diff --git a/kun/src/adapters/tool/builtin-repo-map-tool.test.ts b/kun/src/adapters/tool/builtin-repo-map-tool.test.ts new file mode 100644 index 000000000..ba4ffd2d3 --- /dev/null +++ b/kun/src/adapters/tool/builtin-repo-map-tool.test.ts @@ -0,0 +1,98 @@ +import { mkdtemp, mkdir, rm, writeFile } from 'node:fs/promises' +import { tmpdir } from 'node:os' +import { join } from 'node:path' +import { afterEach, describe, expect, it, vi } from 'vitest' +import { SUBAGENT_READ_ONLY_TOOL_NAMES } from '../../contracts/capabilities.js' +import type { ToolHostContext } from '../../ports/tool-host.js' +import { createRepoMapLocalTool } from './builtin-repo-map-tool.js' + +const tempRoots: string[] = [] + +afterEach(async () => { + await Promise.all(tempRoots.splice(0).map((root) => rm(root, { recursive: true, force: true }))) +}) + +async function createFixture(): Promise { + const root = await mkdtemp(join(tmpdir(), 'kun-repo-map-')) + tempRoots.push(root) + await mkdir(join(root, 'src'), { recursive: true }) + await mkdir(join(root, 'node_modules', 'ignored'), { recursive: true }) + await writeFile(join(root, 'package.json'), JSON.stringify({ name: 'repo-map-fixture' })) + await writeFile( + join(root, 'src', 'user-service.ts'), + [ + "import { openDatabase } from './database'", + 'export class UserRepository {}', + 'export async function findUserByEmail(email: string) {', + ' return openDatabase().users.find((user) => user.email === email)', + '}' + ].join('\n') + ) + await writeFile( + join(root, 'src', 'payment.ts'), + 'export function chargeInvoice(total: number) { return total }\n' + ) + await writeFile( + join(root, 'node_modules', 'ignored', 'user-copy.ts'), + 'export class UserRepositoryCopy {}\n' + ) + return root +} + +function context(workspace: string): ToolHostContext { + return { + threadId: 'thread_repo_map', + turnId: 'turn_repo_map', + workspace, + approvalPolicy: 'auto', + abortSignal: new AbortController().signal, + awaitApproval: vi.fn() + } +} + +describe('repo_map tool', () => { + it('ranks relevant symbols and skips dependency directories', async () => { + const root = await createFixture() + const tool = createRepoMapLocalTool() + + const result = await tool.execute( + { query: 'find user repository by email', maxFiles: 5 }, + context(root) + ) + const output = result.output as { + cache: { hit: boolean } + totals: { scanBackend: string; scannedFiles: number; indexedFiles: number; truncated: boolean } + files: Array<{ + relative_path: string + symbols: Array<{ name: string; kind: string; line: number }> + imports: string[] + }> + skippedDirectories: string[] + } + + expect(output.cache.hit).toBe(false) + expect(output.totals).toMatchObject({ + scanBackend: 'filesystem', + scannedFiles: 3, + indexedFiles: 3, + truncated: false + }) + expect(output.files[0]?.relative_path).toBe('src/user-service.ts') + expect(output.files[0]?.symbols).toEqual(expect.arrayContaining([ + { name: 'UserRepository', kind: 'class', line: 2 }, + { name: 'findUserByEmail', kind: 'function', line: 3 } + ])) + expect(output.files[0]?.imports).toContain('./database') + expect(output.skippedDirectories).toContain('node_modules') + }) + + it('reuses the short-lived index and is enabled for read-only subagents', async () => { + const root = await createFixture() + const tool = createRepoMapLocalTool() + await tool.execute({ query: 'payment' }, context(root)) + + const cached = await tool.execute({ query: 'user' }, context(root)) + expect((cached.output as { cache: { hit: boolean } }).cache.hit).toBe(true) + expect(SUBAGENT_READ_ONLY_TOOL_NAMES).toContain('repo_map') + }) +}) diff --git a/kun/src/adapters/tool/builtin-repo-map-tool.ts b/kun/src/adapters/tool/builtin-repo-map-tool.ts new file mode 100644 index 000000000..84cbf3a08 --- /dev/null +++ b/kun/src/adapters/tool/builtin-repo-map-tool.ts @@ -0,0 +1,721 @@ +import type { Dirent, Stats } from 'node:fs' +import { readFile, readdir, stat } from 'node:fs/promises' +import { basename, extname, join, relative } from 'node:path' +import type { LocalTool } from './local-tool-host.js' +import { + isBinaryBuffer, + normalizeBoolean, + normalizePositiveInteger, + normalizeToolPath, + resolveWorkspacePath, + spawnCapture, + withToolBoundary +} from './builtin-tool-utils.js' + +const DEFAULT_REPO_MAP_MAX_FILES = 20 +const DEFAULT_REPO_MAP_MAX_SYMBOLS = 12 +const DEFAULT_REPO_MAP_SCAN_LIMIT = 2500 +const REPO_MAP_CACHE_TTL_MS = 30_000 +const MAX_SYMBOL_BYTES = 512 * 1024 +const MAX_GIT_RECENT_FILES = 250 +const INDEX_CONCURRENCY = 12 + +const SKIP_DIRS = new Set([ + '.git', + '.hg', + '.svn', + 'node_modules', + '.next', + '.nuxt', + '.turbo', + '.cache', + '.codex', + 'dist', + 'out', + 'build', + 'coverage', + 'target', + '.venv', + 'venv', + '__pycache__', + '.pytest_cache' +]) + +const SOURCE_EXTENSIONS = new Set([ + '.ts', '.tsx', '.js', '.jsx', '.mjs', '.cjs', + '.py', '.rs', '.go', '.java', '.kt', '.kts', + '.cs', '.cpp', '.cc', '.cxx', '.c', '.h', '.hpp', + '.swift', '.rb', '.php', '.vue', '.svelte', + '.json', '.toml', '.yaml', '.yml', '.md', '.mdx' +]) + +const IMPORTANT_FILE_NAMES = new Set([ + 'package.json', + 'tsconfig.json', + 'vite.config.ts', + 'electron.vite.config.ts', + 'README.md', + 'AGENTS.md', + 'CLAUDE.md', + 'pyproject.toml', + 'Cargo.toml', + 'go.mod' +]) + +type RepoMapSymbol = { + name: string + kind: string + line: number +} + +type RepoMapFile = { + path: string + relativePath: string + language: string + size: number + symbols: RepoMapSymbol[] + imports: string[] + tokens: string[] +} + +type RepoMapIndex = { + root: string + target: string + scanBackend: 'git' | 'filesystem' + scannedAt: number + gitHead?: string + files: RepoMapFile[] + scannedFiles: number + skippedDirectories: string[] + truncated: boolean + recentFiles: Map +} + +type RepoMapCacheEntry = { + index: RepoMapIndex + expiresAt: number + scanLimit: number +} + +const repoMapCache = new Map() + +export function createRepoMapLocalTool(): LocalTool { + return { + name: 'repo_map', + description: + 'Build a compact, ranked map of the local codebase before reading files. ' + + 'Uses path/symbol/import extraction, git recency, and BM25-like scoring with a scan budget. ' + + 'Prefer this before broad grep/read passes when you need to understand an unfamiliar repository.', + inputSchema: { + type: 'object', + properties: { + query: { + type: 'string', + description: 'Optional task/search intent used to rank files and symbols.' + }, + path: { + type: 'string', + description: 'Workspace-relative directory or file to map. Defaults to the workspace root.' + }, + maxFiles: { + type: 'number', + description: 'Maximum ranked files to return. Defaults to 20.' + }, + maxSymbolsPerFile: { + type: 'number', + description: 'Maximum symbols returned per file. Defaults to 12.' + }, + maxScanFiles: { + type: 'number', + description: 'Maximum source/config files scanned before truncating. Defaults to 2500.' + }, + refresh: { + type: 'boolean', + description: 'Bypass the short-lived in-process cache.' + } + }, + additionalProperties: false + }, + policy: 'auto', + toolKind: 'tool_call', + execute: async (args, context) => withToolBoundary(async () => { + const rawPath = typeof args.path === 'string' && args.path.trim() ? args.path.trim() : '.' + const query = typeof args.query === 'string' ? args.query.trim() : '' + const maxFiles = clamp(normalizePositiveInteger(args.maxFiles, DEFAULT_REPO_MAP_MAX_FILES), 1, 80) + const maxSymbolsPerFile = clamp(normalizePositiveInteger(args.maxSymbolsPerFile, DEFAULT_REPO_MAP_MAX_SYMBOLS), 0, 50) + const maxScanFiles = clamp(normalizePositiveInteger(args.maxScanFiles, DEFAULT_REPO_MAP_SCAN_LIMIT), 100, 20_000) + const refresh = normalizeBoolean(args.refresh) + const { workspaceRoot, absolutePath, relativePath } = await resolveWorkspacePath(rawPath, context) + const cacheKey = `${workspaceRoot}\0${absolutePath}` + const gitHead = await gitHeadForWorkspace(workspaceRoot) + const cached = repoMapCache.get(cacheKey) + const now = Date.now() + const cacheHit = Boolean( + cached && + !refresh && + cached.expiresAt > now && + cached.scanLimit >= maxScanFiles && + cached.index.gitHead === gitHead + ) + const index = cacheHit + ? cached!.index + : await buildRepoMapIndex({ + workspaceRoot, + target: absolutePath, + maxScanFiles, + gitHead, + signal: context.abortSignal + }) + if (!cacheHit) { + repoMapCache.set(cacheKey, { + index, + expiresAt: now + REPO_MAP_CACHE_TTL_MS, + scanLimit: maxScanFiles + }) + } + + const ranked = rankRepoMapFiles(index.files, query, index.recentFiles) + .slice(0, maxFiles) + .map((entry) => formatRepoMapFile(entry.file, entry.score, entry.reasons, maxSymbolsPerFile)) + + return { + output: { + workspaceRoot, + path: absolutePath, + relative_path: relativePath, + query, + cache: { + hit: cacheHit, + ttlMs: REPO_MAP_CACHE_TTL_MS, + scannedAt: new Date(index.scannedAt).toISOString(), + gitHead: index.gitHead ?? null + }, + totals: { + scanBackend: index.scanBackend, + scannedFiles: index.scannedFiles, + indexedFiles: index.files.length, + truncated: index.truncated + }, + languages: topEntries(countBy(index.files, (file) => file.language), 12), + importantDirectories: topEntries(countBy(index.files, (file) => firstDirectory(file.relativePath)), 12), + entrypoints: entrypoints(index.files).slice(0, 20), + files: ranked, + skippedDirectories: index.skippedDirectories.slice(0, 30), + suggestions: repoMapSuggestions(query, ranked.length, index.truncated) + } + } + }) + } +} + +type BuildRepoMapIndexInput = { + workspaceRoot: string + target: string + maxScanFiles: number + gitHead?: string + signal?: AbortSignal +} + +async function buildRepoMapIndex(input: BuildRepoMapIndexInput): Promise { + const startedAt = Date.now() + const targetStat = await stat(input.target) + const gitDiscovery = targetStat.isDirectory() + ? await discoverGitFiles(input) + : null + const discovery = gitDiscovery ?? await discoverFilesystemFiles(input, targetStat) + const filePaths = discovery.filePaths + const skippedDirectories = discovery.skippedDirectories + const truncated = discovery.truncated + const scanBackend = gitDiscovery ? 'git' : 'filesystem' + const [files, recentFiles] = await Promise.all([ + indexRepoFiles(input.workspaceRoot, filePaths, input.signal), + gitRecentFiles(input.workspaceRoot, input.signal) + ]) + return { + root: input.workspaceRoot, + target: input.target, + scanBackend, + scannedAt: startedAt, + ...(input.gitHead ? { gitHead: input.gitHead } : {}), + files, + scannedFiles: filePaths.length, + skippedDirectories, + truncated, + recentFiles + } +} + +type RepoFileDiscovery = { + filePaths: string[] + skippedDirectories: string[] + truncated: boolean +} + +async function discoverGitFiles(input: BuildRepoMapIndexInput): Promise { + const targetRelative = normalizeToolPath(relative(input.workspaceRoot, input.target) || '.') + try { + const result = await spawnCapture( + 'git', + ['ls-files', '--cached', '--others', '--exclude-standard', '-z', '--', targetRelative], + { cwd: input.workspaceRoot, signal: input.signal } + ) + if (result.exitCode !== 0) return null + const filePaths: string[] = [] + let truncated = false + for (const rawPath of result.stdout.split('\0')) { + const relativePath = normalizeToolPath(rawPath.trim()) + if (!relativePath || shouldSkipRelativePath(relativePath)) continue + const absolutePath = join(input.workspaceRoot, relativePath) + if (!shouldIndexFile(absolutePath)) continue + if (filePaths.length >= input.maxScanFiles) { + truncated = true + break + } + filePaths.push(absolutePath) + } + return { + filePaths, + skippedDirectories: ['gitignored and nested repositories'], + truncated + } + } catch { + return null + } +} + +async function discoverFilesystemFiles( + input: BuildRepoMapIndexInput, + targetStat: Stats +): Promise { + const filePaths: string[] = [] + const skippedDirectories: string[] = [] + const queue = targetStat.isDirectory() ? [input.target] : [] + let queueIndex = 0 + let truncated = false + if (targetStat.isFile() && shouldIndexFile(input.target)) { + filePaths.push(input.target) + } + while (queueIndex < queue.length) { + if (input.signal?.aborted) throw new Error('repo_map aborted') + const current = queue[queueIndex] + queueIndex += 1 + if (!current) break + let entries: Dirent[] + try { + entries = await readdir(current, { withFileTypes: true }) + } catch { + continue + } + if (current !== input.target && entries.some((entry) => entry.isDirectory() && entry.name === '.git')) { + skippedDirectories.push(normalizeToolPath(relative(input.workspaceRoot, current))) + continue + } + entries.sort((a, b) => a.name.localeCompare(b.name)) + for (const entry of entries) { + if (input.signal?.aborted) throw new Error('repo_map aborted') + const next = join(current, entry.name) + if (entry.isDirectory()) { + if (SKIP_DIRS.has(entry.name)) { + skippedDirectories.push(normalizeToolPath(relative(input.workspaceRoot, next) || entry.name)) + continue + } + queue.push(next) + continue + } + if (!entry.isFile() || !shouldIndexFile(next)) continue + if (filePaths.length >= input.maxScanFiles) { + truncated = true + queueIndex = queue.length + break + } + filePaths.push(next) + } + } + return { + filePaths, + skippedDirectories, + truncated + } +} + +function shouldSkipRelativePath(relativePath: string): boolean { + return normalizeToolPath(relativePath) + .split('/') + .some((segment) => SKIP_DIRS.has(segment)) +} + +async function indexRepoFiles( + workspaceRoot: string, + filePaths: string[], + signal?: AbortSignal +): Promise { + const files: RepoMapFile[] = [] + for (let start = 0; start < filePaths.length; start += INDEX_CONCURRENCY) { + if (signal?.aborted) throw new Error('repo_map aborted') + const indexed = await Promise.all( + filePaths.slice(start, start + INDEX_CONCURRENCY).map((filePath) => indexFile(workspaceRoot, filePath)) + ) + for (const file of indexed) { + if (file) files.push(file) + } + } + return files +} + +async function indexFile(workspaceRoot: string, filePath: string): Promise { + let fileStat: Awaited> + try { + fileStat = await stat(filePath) + } catch { + return null + } + if (fileStat.size > 5 * 1024 * 1024) return null + let text = '' + if (fileStat.size <= MAX_SYMBOL_BYTES) { + const buffer = await readFile(filePath) + if (isBinaryBuffer(buffer)) return null + text = buffer.toString('utf8') + } + const relativePath = normalizeToolPath(relative(workspaceRoot, filePath) || basename(filePath)) + const symbols = text ? extractSymbols(relativePath, text) : [] + const imports = text ? extractImports(text) : [] + const language = languageForPath(filePath) + const tokens = tokenize(`${relativePath} ${language} ${symbols.map((symbol) => symbol.name).join(' ')} ${imports.join(' ')}`) + return { + path: filePath, + relativePath, + language, + size: fileStat.size, + symbols, + imports, + tokens + } +} + +function shouldIndexFile(filePath: string): boolean { + const name = basename(filePath) + return IMPORTANT_FILE_NAMES.has(name) || SOURCE_EXTENSIONS.has(extname(filePath).toLowerCase()) +} + +function extractSymbols(relativePath: string, text: string): RepoMapSymbol[] { + const ext = extname(relativePath).toLowerCase() + const out: RepoMapSymbol[] = [] + const lines = text.replace(/\r\n/g, '\n').split('\n') + for (let index = 0; index < lines.length; index += 1) { + const line = lines[index] ?? '' + const symbol = symbolFromLine(ext, line) + if (!symbol) continue + out.push({ ...symbol, line: index + 1 }) + if (out.length >= 120) break + } + return out +} + +function symbolFromLine(ext: string, line: string): Omit | null { + const trimmed = line.trim() + if (!trimmed || trimmed.startsWith('//') || trimmed.startsWith('*') || trimmed.startsWith('#')) return null + const patterns: Array<[RegExp, string]> = [ + [/^(?:export\s+)?(?:default\s+)?(?:async\s+)?function\s+([A-Za-z_$][\w$]*)/, 'function'], + [/^(?:export\s+)?(?:abstract\s+)?class\s+([A-Za-z_$][\w$]*)/, 'class'], + [/^(?:export\s+)?interface\s+([A-Za-z_$][\w$]*)/, 'interface'], + [/^(?:export\s+)?type\s+([A-Za-z_$][\w$]*)/, 'type'], + [/^(?:export\s+)?enum\s+([A-Za-z_$][\w$]*)/, 'enum'], + [/^(?:export\s+)?const\s+([A-Za-z_$][\w$]*)\s*=\s*(?:async\s*)?(?:\([^)]*\)|[A-Za-z_$][\w$]*)\s*=>/, 'function'], + [/^(?:async\s+)?def\s+([A-Za-z_]\w*)\s*\(/, 'function'], + [/^class\s+([A-Za-z_]\w*)/, 'class'], + [/^(?:pub\s+)?(?:async\s+)?fn\s+([A-Za-z_]\w*)/, 'function'], + [/^(?:pub\s+)?(?:struct|enum|trait)\s+([A-Za-z_]\w*)/, 'type'], + [/^func\s+(?:\([^)]+\)\s*)?([A-Za-z_]\w*)\s*\(/, 'function'], + [/^type\s+([A-Za-z_]\w*)\s+/, 'type'], + [/^(?:public\s+|private\s+|protected\s+|static\s+|final\s+|abstract\s+)*class\s+([A-Za-z_]\w*)/, 'class'], + [/^(?:public\s+|private\s+|protected\s+|static\s+|final\s+|abstract\s+)*(?:interface|enum)\s+([A-Za-z_]\w*)/, 'type'] + ] + for (const [pattern, kind] of patterns) { + const match = trimmed.match(pattern) + if (match?.[1]) return { name: match[1], kind } + } + if (ext === '.rs') { + const impl = trimmed.match(/^impl(?:<[^>]+>)?\s+([A-Za-z_]\w*)/) + if (impl?.[1]) return { name: impl[1], kind: 'impl' } + } + return null +} + +function extractImports(text: string): string[] { + const out: string[] = [] + const seen = new Set() + for (const line of text.replace(/\r\n/g, '\n').split('\n')) { + const trimmed = line.trim() + const match = + trimmed.match(/^import\s+.*?\s+from\s+['"]([^'"]+)['"]/) ?? + trimmed.match(/^import\s+['"]([^'"]+)['"]/) ?? + trimmed.match(/^export\s+.*?\s+from\s+['"]([^'"]+)['"]/) ?? + trimmed.match(/^const\s+\w+\s*=\s*require\(['"]([^'"]+)['"]\)/) ?? + trimmed.match(/^from\s+([A-Za-z0-9_.]+)\s+import\s+/) ?? + trimmed.match(/^use\s+([A-Za-z0-9_:]+)/) ?? + trimmed.match(/^package\s+([A-Za-z0-9_.]+)/) + const value = match?.[1]?.trim() + if (!value || seen.has(value)) continue + seen.add(value) + out.push(value) + if (out.length >= 30) break + } + return out +} + +function rankRepoMapFiles( + files: RepoMapFile[], + query: string, + recentFiles: Map +): Array<{ file: RepoMapFile; score: number; reasons: string[] }> { + const queryTokens = tokenize(query) + const documentFrequency = documentFrequencyFor(files) + const averageLength = files.length > 0 + ? Math.max(1, files.reduce((total, file) => total + file.tokens.length, 0) / files.length) + : 1 + return files + .map((file) => { + const reasons: string[] = [] + const queryScore = queryTokens.length > 0 + ? bm25(file.tokens, queryTokens, documentFrequency, files.length, averageLength) + : 0 + const rawPathBoost = importantPathBoost(file.relativePath) + const pathBoost = queryTokens.length > 0 ? rawPathBoost * 0.25 : rawPathBoost + const recency = recentFiles.get(file.relativePath) ?? 0 + if (queryScore > 0) reasons.push('query_match') + if (pathBoost > 0) reasons.push('entrypoint_or_core_path') + if (recency > 0) reasons.push('git_recent') + if (file.symbols.length > 0) reasons.push('symbols') + return { + file, + score: Number((queryScore + pathBoost + recency).toFixed(3)), + reasons + } + }) + .sort((a, b) => b.score - a.score || a.file.relativePath.localeCompare(b.file.relativePath)) +} + +function bm25( + documentTokens: string[], + queryTokens: string[], + documentFrequency: Map, + documentCount: number, + averageLength: number +): number { + const tf = new Map() + for (const token of documentTokens) tf.set(token, (tf.get(token) ?? 0) + 1) + let score = 0 + const k1 = 1.2 + const b = 0.75 + const docLength = Math.max(documentTokens.length, 1) + for (const token of new Set(queryTokens)) { + const frequency = tf.get(token) ?? 0 + if (frequency === 0) continue + const df = documentFrequency.get(token) ?? 0 + const idf = Math.log(1 + (documentCount - df + 0.5) / (df + 0.5)) + const normalized = (frequency * (k1 + 1)) / (frequency + k1 * (1 - b + b * (docLength / averageLength))) + score += idf * normalized + } + return score +} + +function documentFrequencyFor(files: RepoMapFile[]): Map { + const out = new Map() + for (const file of files) { + for (const token of new Set(file.tokens)) { + out.set(token, (out.get(token) ?? 0) + 1) + } + } + return out +} + +function formatRepoMapFile( + file: RepoMapFile, + score: number, + reasons: string[], + maxSymbolsPerFile: number +): Record { + return { + path: file.path, + relative_path: file.relativePath, + language: file.language, + size: file.size, + score, + reasons, + symbols: file.symbols.slice(0, maxSymbolsPerFile), + symbolCount: file.symbols.length, + imports: file.imports.slice(0, 12) + } +} + +function entrypoints(files: RepoMapFile[]): string[] { + return files + .filter((file) => importantPathBoost(file.relativePath) > 0) + .sort((a, b) => importantPathBoost(b.relativePath) - importantPathBoost(a.relativePath)) + .map((file) => file.relativePath) +} + +function importantPathBoost(relativePath: string): number { + const name = basename(relativePath) + if (IMPORTANT_FILE_NAMES.has(name)) return 2.5 + if (/^(src|app|packages|kun)\//.test(relativePath)) return 0.8 + if (/(^|\/)(index|main|runtime|server|router|store|config)\.[^.]+$/.test(relativePath)) return 1.4 + return 0 +} + +async function gitHeadForWorkspace(workspaceRoot: string): Promise { + try { + const result = await spawnCapture('git', ['rev-parse', 'HEAD'], { cwd: workspaceRoot }) + const value = result.stdout.trim() + return result.exitCode === 0 && value ? value : undefined + } catch { + return undefined + } +} + +async function gitRecentFiles(workspaceRoot: string, signal?: AbortSignal): Promise> { + const out = new Map() + await addGitStatusRecency(workspaceRoot, out, signal) + await addGitLogRecency(workspaceRoot, out, signal) + return out +} + +async function addGitStatusRecency( + workspaceRoot: string, + out: Map, + signal?: AbortSignal +): Promise { + try { + const result = await spawnCapture('git', ['status', '--short'], { cwd: workspaceRoot, signal }) + if (result.exitCode !== 0) return + for (const line of result.stdout.split(/\r?\n/)) { + const rawPath = line.slice(3).trim().replace(/^"|"$/g, '') + if (!rawPath) continue + const normalized = normalizeToolPath(rawPath.includes(' -> ') ? rawPath.split(' -> ').pop() ?? rawPath : rawPath) + out.set(normalized, Math.max(out.get(normalized) ?? 0, 3)) + if (out.size >= MAX_GIT_RECENT_FILES) return + } + } catch { + // Git metadata is a boost, not a requirement. + } +} + +async function addGitLogRecency( + workspaceRoot: string, + out: Map, + signal?: AbortSignal +): Promise { + try { + const result = await spawnCapture( + 'git', + ['log', '--name-only', '--pretty=format:', '-n', '30'], + { cwd: workspaceRoot, signal } + ) + if (result.exitCode !== 0) return + let rank = 0 + for (const line of result.stdout.split(/\r?\n/)) { + const file = normalizeToolPath(line.trim()) + if (!file || out.has(file)) continue + const score = Math.max(0.2, 1.5 - rank * 0.02) + out.set(file, score) + rank += 1 + if (out.size >= MAX_GIT_RECENT_FILES) return + } + } catch { + // Git metadata is a boost, not a requirement. + } +} + +function tokenize(text: string): string[] { + const tokens: string[] = [] + const normalized = text + .normalize('NFKC') + .replace(/([a-z0-9])([A-Z])/g, '$1 $2') + .toLowerCase() + for (const match of normalized.matchAll(/[a-z0-9][a-z0-9_-]{1,}/g)) { + const token = match[0] + for (const part of token.split(/[_-]+/)) { + if (part.length >= 2 && !/^\d+$/.test(part)) tokens.push(part) + } + if (token.length >= 2 && !/^\d+$/.test(token)) tokens.push(token) + } + for (const segment of normalized.match(/\p{Script=Han}+/gu) ?? []) { + const chars = [...segment] + for (let size = 2; size <= Math.min(4, chars.length); size += 1) { + for (let index = 0; index <= chars.length - size; index += 1) { + tokens.push(chars.slice(index, index + size).join('')) + } + } + } + return tokens +} + +function countBy(values: T[], keyFor: (value: T) => string): Map { + const out = new Map() + for (const value of values) { + const key = keyFor(value) + if (!key) continue + out.set(key, (out.get(key) ?? 0) + 1) + } + return out +} + +function topEntries(map: Map, limit: number): Array<{ name: string; count: number }> { + return [...map.entries()] + .sort((a, b) => b[1] - a[1] || a[0].localeCompare(b[0])) + .slice(0, limit) + .map(([name, count]) => ({ name, count })) +} + +function firstDirectory(relativePath: string): string { + const normalized = normalizeToolPath(relativePath) + const slash = normalized.indexOf('/') + return slash === -1 ? '.' : normalized.slice(0, slash) +} + +function languageForPath(filePath: string): string { + const name = basename(filePath) + if (IMPORTANT_FILE_NAMES.has(name)) return name + const ext = extname(filePath).toLowerCase() + switch (ext) { + case '.ts': + case '.tsx': + return 'typescript' + case '.js': + case '.jsx': + case '.mjs': + case '.cjs': + return 'javascript' + case '.py': + return 'python' + case '.rs': + return 'rust' + case '.go': + return 'go' + case '.md': + case '.mdx': + return 'markdown' + case '.json': + return 'json' + case '.yaml': + case '.yml': + return 'yaml' + default: + return ext.replace(/^\./, '') || 'text' + } +} + +function repoMapSuggestions(query: string, resultCount: number, truncated: boolean): string[] { + const out = [ + resultCount > 0 + ? 'Read the highest-ranked files before editing; use grep/lsp for exact call sites.' + : 'No source files matched the current scope; try a broader path or refresh=true.' + ] + if (!query) out.push('Pass query to rank files for the current task instead of returning only entrypoint/core-path scores.') + if (truncated) out.push('The scan hit maxScanFiles; narrow path or increase maxScanFiles for a fuller map.') + return out +} + +function clamp(value: number, min: number, max: number): number { + return Math.max(min, Math.min(max, value)) +} diff --git a/kun/src/adapters/tool/builtin-tool-types.ts b/kun/src/adapters/tool/builtin-tool-types.ts index 24f0f1873..bbdeceb51 100644 --- a/kun/src/adapters/tool/builtin-tool-types.ts +++ b/kun/src/adapters/tool/builtin-tool-types.ts @@ -98,6 +98,7 @@ export type BuiltinToolName = | 'find' | 'ls' | 'lsp' + | 'repo_map' | 'verify_changes' export const allBuiltinToolNames: Set = new Set([ 'read', @@ -108,6 +109,7 @@ export const allBuiltinToolNames: Set = new Set([ 'find', 'ls', 'lsp', + 'repo_map', 'verify_changes' ]) export type ToolName = BuiltinToolName diff --git a/kun/src/adapters/tool/builtin-tools.ts b/kun/src/adapters/tool/builtin-tools.ts index 77108b4e7..f25ac2e4e 100644 --- a/kun/src/adapters/tool/builtin-tools.ts +++ b/kun/src/adapters/tool/builtin-tools.ts @@ -10,6 +10,7 @@ import { createEditLocalTool, createWriteLocalTool } from './builtin-file-tools. import { createLspLocalTool } from './builtin-lsp-tool.js' import { createReadLocalTool } from './builtin-read-tool.js' import { createFindLocalTool, createGrepLocalTool, createLsLocalTool } from './builtin-search-tools.js' +import { createRepoMapLocalTool } from './builtin-repo-map-tool.js' import { createVerifyChangesLocalTool } from './builtin-verify-tool.js' export * from './builtin-tool-types.js' @@ -17,6 +18,7 @@ export * from './builtin-tool-operations.js' export * from './builtin-read-tool.js' export * from './builtin-file-tools.js' export * from './builtin-search-tools.js' +export * from './builtin-repo-map-tool.js' export * from './builtin-bash-tool.js' export * from './builtin-verify-tool.js' @@ -41,6 +43,8 @@ export function createBuiltinLocalTool( return createLsLocalTool(options.ls) case 'lsp': return createLspLocalTool() + case 'repo_map': + return createRepoMapLocalTool() case 'verify_changes': return createVerifyChangesLocalTool() } @@ -64,6 +68,7 @@ export function buildBuiltinLocalTools(options: BuiltinLocalToolsOptions = {}): createFindLocalTool(options.find), createLsLocalTool(options.ls), createLspLocalTool(), + createRepoMapLocalTool(), createVerifyChangesLocalTool() ] } @@ -90,7 +95,8 @@ export function buildReadOnlyBuiltinLocalTools(options: BuiltinLocalToolsOptions createReadLocalTool(options.read), createGrepLocalTool(options.grep), createFindLocalTool(options.find), - createLsLocalTool(options.ls) + createLsLocalTool(options.ls), + createRepoMapLocalTool() ] } @@ -110,6 +116,7 @@ export function buildBuiltinLocalToolRecord( find: createFindLocalTool(options.find), ls: createLsLocalTool(options.ls), lsp: createLspLocalTool(), + repo_map: createRepoMapLocalTool(), verify_changes: createVerifyChangesLocalTool() } } diff --git a/kun/src/adapters/tool/capability-registry.ts b/kun/src/adapters/tool/capability-registry.ts index ecabc78a4..81694a2d3 100644 --- a/kun/src/adapters/tool/capability-registry.ts +++ b/kun/src/adapters/tool/capability-registry.ts @@ -29,6 +29,7 @@ const PLAN_MODE_ALLOWED_TOOL_NAMES = new Set([ 'grep', 'find', 'ls', + 'repo_map', 'create_plan', 'user_input', 'request_user_input' diff --git a/kun/src/adapters/tool/delegation-tool-provider.ts b/kun/src/adapters/tool/delegation-tool-provider.ts index 0a4ff3083..889fcf928 100644 --- a/kun/src/adapters/tool/delegation-tool-provider.ts +++ b/kun/src/adapters/tool/delegation-tool-provider.ts @@ -30,6 +30,21 @@ export function buildDelegationToolProviders(runtime: DelegationRuntime | undefi detach: { type: 'boolean', description: 'Fire-and-forget. The call returns immediately with a queued/running record; the child keeps executing in the background and can be checked via diagnostics or aborted from the GUI.' + }, + tokenBudget: { + type: 'integer', + minimum: 1, + description: 'Optional hard cap for total child tokens.' + }, + timeBudgetMs: { + type: 'integer', + minimum: 1, + description: 'Optional wall-clock timeout in milliseconds.' + }, + returnFormat: { + type: 'string', + enum: ['summary', 'evidence'], + description: 'Require either a normal summary or explicit evidence items.' } }, required: ['prompt'], @@ -39,6 +54,12 @@ export function buildDelegationToolProviders(runtime: DelegationRuntime | undefi execute: async (args, context, onUpdate) => { const prompt = typeof args.prompt === 'string' ? args.prompt.trim() : '' if (!prompt) return { output: { error: 'prompt is required' }, isError: true } + if (args.tokenBudget !== undefined && !isPositiveInteger(args.tokenBudget)) { + return { output: { error: 'tokenBudget must be a positive integer' }, isError: true } + } + if (args.timeBudgetMs !== undefined && !isPositiveInteger(args.timeBudgetMs)) { + return { output: { error: 'timeBudgetMs must be a positive integer' }, isError: true } + } const record = await runtime.runChild({ parentThreadId: context.threadId, parentTurnId: context.turnId, @@ -48,6 +69,9 @@ export function buildDelegationToolProviders(runtime: DelegationRuntime | undefi ...(typeof args.model === 'string' ? { model: args.model } : {}), ...(typeof args.profile === 'string' ? { profile: args.profile } : {}), ...(args.detach === true ? { detach: true } : {}), + ...(isPositiveInteger(args.tokenBudget) ? { tokenBudget: args.tokenBudget } : {}), + ...(isPositiveInteger(args.timeBudgetMs) ? { timeBudgetMs: args.timeBudgetMs } : {}), + ...(args.returnFormat === 'evidence' ? { returnFormat: 'evidence' as const } : {}), // Emit a partial result the moment the child id exists, so the GUI // can offer "open session" (and stream the child live) while the // child is still running — not only after it completes. @@ -65,7 +89,12 @@ export function buildDelegationToolProviders(runtime: DelegationRuntime | undefi status: record.status, summary: record.summary, error: record.error, + evidence: record.evidence, usage: record.usage, + returnFormat: record.returnFormat, + ...(record.tokenBudget ? { tokenBudget: record.tokenBudget } : {}), + ...(record.timeBudgetMs ? { timeBudgetMs: record.timeBudgetMs } : {}), + ...(record.budgetExceeded ? { budgetExceeded: record.budgetExceeded } : {}), ...(record.profile ? { profile: record.profile } : {}), ...(record.toolPolicy ? { toolPolicy: record.toolPolicy } : {}), ...(record.toolInvocations !== undefined ? { toolInvocations: record.toolInvocations } : {}), @@ -80,6 +109,10 @@ export function buildDelegationToolProviders(runtime: DelegationRuntime | undefi }] } +function isPositiveInteger(value: unknown): value is number { + return typeof value === 'number' && Number.isInteger(value) && value > 0 +} + function buildDelegateTaskDescription( runtime: DelegationRuntime, profiles: { name: string; mode: string; toolPolicy: string; model?: string; providerId?: string; description?: string }[] diff --git a/kun/src/adapters/tool/local-tool-host.test.ts b/kun/src/adapters/tool/local-tool-host.test.ts index adfe4f2b6..918277bf4 100644 --- a/kun/src/adapters/tool/local-tool-host.test.ts +++ b/kun/src/adapters/tool/local-tool-host.test.ts @@ -1,6 +1,7 @@ import { describe, expect, it, vi } from 'vitest' import { LocalToolHost, echoTool } from './local-tool-host.js' import type { ToolHostContext } from '../../ports/tool-host.js' +import { InMemoryArtifactStore } from '../../artifacts/artifact-store.js' describe('LocalToolHost approval policy', () => { it('asks before auto tools when approval policy is always', async () => { @@ -26,4 +27,34 @@ describe('LocalToolHost approval policy', () => { expect(awaitApproval).toHaveBeenCalledTimes(1) expect(result.approved).toBe(false) }) + + it('offloads oversized successful tool output to the artifact store', async () => { + const artifactStore = new InMemoryArtifactStore() + const host = new LocalToolHost({ tools: [LocalToolHost.defineTool({ + name: 'large_output', + description: 'returns a large payload', + inputSchema: { type: 'object' }, + execute: async () => ({ output: 'x'.repeat(140 * 1024) }) + })] }) + const result = await host.execute( + { callId: 'call_large', toolName: 'large_output', arguments: {} }, + { + threadId: 'thread_1', + turnId: 'turn_1', + workspace: '/tmp/workspace', + approvalPolicy: 'auto', + sandboxMode: 'danger-full-access', + artifactStore, + abortSignal: new AbortController().signal, + awaitApproval: vi.fn(async () => 'allow' as const) + } + ) + expect(result.item).toMatchObject({ + kind: 'tool_result', + output: { artifactId: expect.stringMatching(/^art_/), truncated: true } + }) + if (result.item.kind !== 'tool_result') throw new Error('expected tool result') + const artifactId = String((result.item.output as Record).artifactId) + expect(await artifactStore.get(artifactId)).toHaveLength(140 * 1024) + }) }) diff --git a/kun/src/adapters/tool/local-tool-host.ts b/kun/src/adapters/tool/local-tool-host.ts index a80070ed8..6066c3e08 100644 --- a/kun/src/adapters/tool/local-tool-host.ts +++ b/kun/src/adapters/tool/local-tool-host.ts @@ -230,7 +230,7 @@ export class LocalToolHost implements ToolHost { } } const rateLimited = normalizeRateLimitedToolOutput(hookedResult.output) - const output = rateLimited.rateLimited ? rateLimited.output : hookedResult.output + let output = rateLimited.rateLimited ? rateLimited.output : hookedResult.output const isError = hookedResult.isError || rateLimited.isError this.readTracker.observeToolResult({ context, @@ -238,6 +238,7 @@ export class LocalToolHost implements ToolHost { output, isError }) + if (!isError) output = await offloadLargeToolOutput(output, activeCall.toolName, context) const item = makeToolResultItem({ id: `item_${activeCall.callId}`, turnId: context.turnId, @@ -340,6 +341,35 @@ export class LocalToolHost implements ToolHost { } } +const ARTIFACT_OUTPUT_THRESHOLD_BYTES = 128 * 1024 + +async function offloadLargeToolOutput( + output: unknown, + toolName: string, + context: ToolHostContext +): Promise { + if (!context.artifactStore) return output + let content: string + try { + content = typeof output === 'string' ? output : JSON.stringify(output) + } catch { + return output + } + if (Buffer.byteLength(content, 'utf8') <= ARTIFACT_OUTPUT_THRESHOLD_BYTES) return output + try { + const stored = await context.artifactStore.put({ content, source: 'tool', origin: toolName }) + return { + artifactId: stored.meta.id, + byteSize: stored.meta.byteSize, + lineCount: stored.meta.lineCount, + truncated: stored.summary.truncated, + preview: stored.summary.inline + } + } catch { + return output + } +} + function hookContext( context: ToolHostContext ): Pick { diff --git a/kun/src/adapters/tool/mcp-naming.ts b/kun/src/adapters/tool/mcp-naming.ts new file mode 100644 index 000000000..ef40edd67 --- /dev/null +++ b/kun/src/adapters/tool/mcp-naming.ts @@ -0,0 +1,90 @@ +import { createHash } from 'node:crypto' +import type { McpServerConfig } from '../../contracts/capabilities.js' + +/** + * MCP identifier, slug, and workspace-trust helpers. + * + * Pure string/path logic shared by the tool, OAuth, and transport layers. + * Kept dependency-free so every other MCP module can import it without + * creating an import cycle. + */ + +export function normalizeMcpToolName(serverId: string, toolName: string): string { + return `mcp_${slug(serverId)}_${slug(toolName)}` +} + +export function isMcpServerTrusted(server: McpServerConfig, workspace: string): boolean { + if (server.trustScope === 'user') return true + return workspaceMatchesRoots(workspace, server.trustedWorkspaceRoots) +} + +export function isMcpServerVisible(server: McpServerConfig, workspace: string): boolean { + if (server.workspaceRoots.length === 0) return true + return workspaceMatchesRoots(workspace, server.workspaceRoots) +} + +export function canUseMcpServer(server: McpServerConfig, workspace: string): boolean { + return isMcpServerVisible(server, workspace) && isMcpServerTrusted(server, workspace) +} + +export function resolveMcpServerCwd(server: McpServerConfig): string | undefined { + if (server.transport !== 'stdio') return undefined + const configured = server.cwd?.trim() + if (configured) return configured + if (server.trustScope !== 'workspace') return undefined + return server.trustedWorkspaceRoots.map((root) => root.trim()).find(Boolean) +} + +export function slug(value: string): string { + let out = '' + for (const char of value.trim().toLowerCase()) { + if (isSlugChar(char)) { + out += char + } else if (out && out[out.length - 1] !== '_') { + out += '_' + } + } + return trimBoundaryUnderscores(out) || 'tool' +} + +export function hashText(value: string): string { + return createHash('sha256').update(value).digest('hex') +} + +export function catalogFingerprint(values: readonly string[]): string { + return createHash('sha256') + .update(JSON.stringify([...values].sort())) + .digest('hex') + .slice(0, 16) +} + +function workspaceMatchesRoots(workspace: string, roots: readonly string[]): boolean { + const normalizedWorkspace = normalizePathForTrust(workspace) + return roots.some((root) => { + const normalizedRoot = normalizePathForTrust(root) + return normalizedWorkspace === normalizedRoot || normalizedWorkspace.startsWith(`${normalizedRoot}/`) + }) +} + +function normalizePathForTrust(value: string): string { + return trimTrailingSlashes(value.replaceAll('\\', '/')) +} + +function isSlugChar(char: string): boolean { + const code = char.charCodeAt(0) + return char === '_' || (code >= 48 && code <= 57) || (code >= 97 && code <= 122) +} + +function trimBoundaryUnderscores(value: string): string { + let start = 0 + let end = value.length + while (start < end && value[start] === '_') start += 1 + while (end > start && value[end - 1] === '_') end -= 1 + return value.slice(start, end) +} + +function trimTrailingSlashes(value: string): string { + let end = value.length + while (end > 0 && value.charCodeAt(end - 1) === 47) end -= 1 + return end === value.length ? value : value.slice(0, end) +} diff --git a/kun/src/adapters/tool/mcp-oauth-provider.ts b/kun/src/adapters/tool/mcp-oauth-provider.ts new file mode 100644 index 000000000..7e9bc229e --- /dev/null +++ b/kun/src/adapters/tool/mcp-oauth-provider.ts @@ -0,0 +1,458 @@ +import { + type OAuthClientProvider, + type OAuthDiscoveryState +} from '@modelcontextprotocol/sdk/client/auth.js' +import type { + OAuthClientInformationMixed, + OAuthClientMetadata, + OAuthTokens +} from '@modelcontextprotocol/sdk/shared/auth.js' +import { spawn } from 'node:child_process' +import { createHash, randomBytes } from 'node:crypto' +import { createServer } from 'node:http' +import { join } from 'node:path' +import type { McpCapabilityConfig, McpServerConfig } from '../../contracts/capabilities.js' +import { hashText, slug } from './mcp-naming.js' +import { FileMcpOAuthStore, type McpOAuthState } from './mcp-oauth-store.js' +import type { SecretEncryptor } from '../../security/secret-store.js' +import { + McpAuthorizationRequiredError, + type McpOAuthClearResult, + type McpOAuthDiagnostic, + type McpOAuthStatus +} from './mcp-types.js' + +const MCP_OAUTH_REDIRECT_HOST = '127.0.0.1' +const MCP_OAUTH_REDIRECT_PATH = '/oauth/callback' +const MCP_OAUTH_PORT_BASE = 49_152 +const MCP_OAUTH_PORT_RANGE = 12_000 + +export type McpOAuthProviderOptions = { + storageDir: string + openExternal?: (url: URL) => void | Promise + /** + * When false (the default) the provider refuses to start the loopback + * callback server or open a browser: `redirectToAuthorization` throws + * {@link McpAuthorizationRequiredError} synchronously. This is the load- + * bearing guard that keeps a non-interactive startup connect from opening a + * browser the moment the SDK hits a 401 — the interactive flag must reach the + * provider, not just the connect wrapper. + */ + interactive?: boolean + /** Optional encryptor; when present, persisted OAuth tokens are encrypted at rest. */ + encryptor?: SecretEncryptor +} + +type McpOAuthDiagnosticDetail = Omit< + McpOAuthDiagnostic, + 'serverId' | 'enabled' | 'configured' | 'transport' | 'url' +> + +/** + * File-backed {@link OAuthClientProvider} for a single remote MCP server. + * + * Responsibilities are deliberately narrow: implement the SDK's OAuth client + * contract, run a one-shot loopback callback server for the authorization + * code, and persist credential material through {@link FileMcpOAuthStore}. + * Transport wiring and connection orchestration live elsewhere so this class + * stays focused on the authorization state machine. + */ +export class FileMcpOAuthProvider implements OAuthClientProvider { + readonly clientMetadataUrl?: string + private readonly store: FileMcpOAuthStore + private pendingAuthorizationCode: Promise | null = null + private pendingState: string | null = null + + constructor( + private readonly serverId: string, + private readonly server: McpServerConfig, + storagePath: string, + private readonly openExternal: (url: URL) => void | Promise = defaultOpenExternal, + private readonly now: () => number = () => Date.now(), + /** + * When false the provider never starts a callback server or opens a + * browser; `redirectToAuthorization` throws immediately. Defaults to false + * so an accidental interactive trigger during startup is impossible. + */ + private readonly interactive: boolean = false, + /** Optional encryptor; when present, persisted tokens are encrypted at rest. */ + encryptor?: SecretEncryptor + ) { + this.store = new FileMcpOAuthStore(storagePath, encryptor) + } + + get redirectUrl(): URL { + return new URL(`http://${MCP_OAUTH_REDIRECT_HOST}:${this.redirectPort()}${MCP_OAUTH_REDIRECT_PATH}`) + } + + get clientMetadata(): OAuthClientMetadata { + return { + client_name: this.server.oauth?.clientName ?? `Kun MCP Client (${this.serverId})`, + redirect_uris: [this.redirectUrl.toString()], + grant_types: ['authorization_code', 'refresh_token'], + response_types: ['code'], + token_endpoint_auth_method: this.server.oauth?.clientSecret ? 'client_secret_post' : 'none', + ...(this.server.oauth?.scopes.length ? { scope: this.server.oauth.scopes.join(' ') } : {}) + } + } + + async clientInformation(): Promise { + const configured = this.configuredClientInformation() + if (configured) return configured + return (await this.store.read()).clientInformation + } + + async saveClientInformation(clientInformation: OAuthClientInformationMixed): Promise { + if (this.configuredClientInformation()) return + await this.store.update((state) => ({ ...state, clientInformation })) + } + + async tokens(): Promise { + return (await this.store.read()).tokens + } + + async saveTokens(tokens: OAuthTokens): Promise { + await this.store.update((state) => { + const next: McpOAuthState = { ...state, tokens, tokensObtainedAt: this.now() } + // A fresh token set means the prior authorization failure (if any) is + // resolved; clear it so diagnostics flip from "error" to "authorized". + delete next.lastError + delete next.lastErrorAt + return next + }) + } + + state(): string { + this.pendingState = randomBytes(24).toString('hex') + return this.pendingState + } + + async redirectToAuthorization(authorizationUrl: URL): Promise { + // Non-interactive providers must never reach the browser/callback. Throw + // BEFORE starting the loopback server or opening a URL, because the SDK + // calls this synchronously inside auth() the moment a 401 is seen — even + // a started-then-aborted callback would briefly bind the redirect port and + // could pop a browser tab on startup. The connection layer maps this to a + // "waiting for user" state instead of a transport error. + if (!this.interactive) { + throw new McpAuthorizationRequiredError(this.serverId) + } + if (authorizationUrl.protocol !== 'http:' && authorizationUrl.protocol !== 'https:') { + throw new Error(`MCP OAuth authorization URL must use http or https for server "${this.serverId}"`) + } + const callback = this.startCallbackServer() + this.pendingAuthorizationCode = callback.code + try { + await callback.ready + await this.openExternal(authorizationUrl) + } catch (error) { + callback.cancel() + this.pendingAuthorizationCode = null + this.pendingState = null + throw error + } + } + + async waitForAuthorizationCode(): Promise { + const authorizationCode = this.pendingAuthorizationCode + if (!authorizationCode) { + throw new Error(`MCP OAuth authorization was not started for server "${this.serverId}"`) + } + try { + return await authorizationCode + } finally { + if (this.pendingAuthorizationCode === authorizationCode) { + this.pendingAuthorizationCode = null + this.pendingState = null + } + } + } + + async saveCodeVerifier(codeVerifier: string): Promise { + await this.store.update((state) => ({ ...state, codeVerifier })) + } + + async codeVerifier(): Promise { + const codeVerifier = (await this.store.read()).codeVerifier + if (!codeVerifier) throw new Error(`MCP OAuth code verifier is missing for server "${this.serverId}"`) + return codeVerifier + } + + async saveDiscoveryState(discoveryState: OAuthDiscoveryState): Promise { + await this.store.update((state) => ({ ...state, discoveryState })) + } + + async discoveryState(): Promise { + return (await this.store.read()).discoveryState + } + + async invalidateCredentials(scope: 'all' | 'client' | 'tokens' | 'verifier' | 'discovery'): Promise { + if (scope === 'all') { + await this.store.clear() + return + } + await this.store.update((state) => { + const next = { ...state } + if (scope === 'client') delete next.clientInformation + if (scope === 'tokens') { + delete next.tokens + delete next.tokensObtainedAt + } + if (scope === 'verifier') delete next.codeVerifier + if (scope === 'discovery') delete next.discoveryState + return next + }) + } + + /** + * Persist the reason an authorization attempt failed. Surfaced by + * diagnostics as status `error` so the GUI can explain a stuck connector + * (callback error, timeout, cancellation) instead of showing a bare state. + */ + async recordAuthorizationError(message: string): Promise { + await this.store.update((state) => ({ + ...state, + lastError: message, + lastErrorAt: new Date(this.now()).toISOString() + })) + } + + async diagnostics(): Promise { + const state = await this.store.read() + const hasClientInformation = Boolean(state.clientInformation) + const hasTokens = Boolean(state.tokens?.access_token) + const hasRefreshToken = Boolean(state.tokens?.refresh_token) + const hasCodeVerifier = Boolean(state.codeVerifier) + const hasDiscoveryState = Boolean(state.discoveryState) + const expiresAt = computeTokenExpiry(state) + const expired = hasTokens && expiresAt !== undefined && Date.parse(expiresAt) <= this.now() + const grantedScopes = parseTokenScopes(state.tokens?.scope) + const status = deriveOAuthStatus({ + hasTokens, + expired, + hasPartialState: hasClientInformation || hasCodeVerifier || hasDiscoveryState, + hasError: Boolean(state.lastError) + }) + return { + status, + hasClientInformation, + hasTokens, + hasRefreshToken, + hasCodeVerifier, + hasDiscoveryState, + ...(grantedScopes.length ? { grantedScopes } : {}), + ...(expiresAt ? { expiresAt } : {}), + ...(state.lastError ? { lastError: state.lastError } : {}), + ...(state.lastErrorAt ? { lastErrorAt: state.lastErrorAt } : {}) + } + } + + async clearCredentials(): Promise { + await this.store.clear() + } + + private configuredClientInformation(): OAuthClientInformationMixed | undefined { + if (!this.server.oauth?.clientId) return undefined + return { + client_id: this.server.oauth.clientId, + ...(this.server.oauth.clientSecret ? { client_secret: this.server.oauth.clientSecret } : {}) + } + } + + private redirectPort(): number { + return this.server.oauth?.redirectPort ?? defaultMcpOAuthRedirectPort(this.serverId, this.server.url ?? '') + } + + private startCallbackServer(): { code: Promise; ready: Promise; cancel: () => void } { + const timeoutMs = this.server.oauth?.callbackTimeoutMs ?? 120_000 + let resolveCode!: (code: string) => void + let rejectCode!: (error: Error) => void + let completed = false + let listening = false + let timer: ReturnType | undefined + const codePromise = new Promise((resolve, reject) => { + resolveCode = resolve + rejectCode = reject + }) + const closeServer = () => { + if (timer) clearTimeout(timer) + if (listening) { + server.close() + listening = false + } + } + const resolveOnce = (code: string) => { + if (completed) return + completed = true + closeServer() + resolveCode(code) + } + const rejectOnce = (error: Error) => { + if (completed) return + completed = true + closeServer() + rejectCode(error) + } + const server = createServer((request, response) => { + const url = new URL(request.url ?? '/', this.redirectUrl) + if (url.pathname !== MCP_OAUTH_REDIRECT_PATH) { + response.writeHead(404, { 'content-type': 'text/plain; charset=utf-8' }) + response.end('Not found') + return + } + const state = url.searchParams.get('state') + if (!this.pendingState || state !== this.pendingState) { + response.writeHead(400, { 'content-type': 'text/html; charset=utf-8' }) + response.end('Kun OAuth

Authorization state mismatch.

') + return + } + const code = url.searchParams.get('code') + const error = url.searchParams.get('error') + if (code) { + response.writeHead(200, { 'content-type': 'text/html; charset=utf-8' }) + response.end('Kun OAuth

Authorization complete. You can close this window.

') + resolveOnce(code) + } else { + response.writeHead(400, { 'content-type': 'text/html; charset=utf-8' }) + response.end('Kun OAuth

Authorization failed.

') + rejectOnce(new Error(error ? `MCP OAuth authorization failed: ${error}` : 'MCP OAuth callback did not include a code')) + } + }) + timer = setTimeout(() => { + rejectOnce(new Error(`MCP OAuth authorization timed out for server "${this.serverId}"`)) + }, timeoutMs) + const ready = new Promise((resolve, reject) => { + server.once('error', (error) => { + rejectOnce(error instanceof Error ? error : new Error(String(error))) + reject(error) + }) + server.listen(this.redirectPort(), MCP_OAUTH_REDIRECT_HOST, () => { + listening = true + resolve() + }) + }) + timer.unref() + codePromise.finally(() => { + if (timer) clearTimeout(timer) + }).catch(() => undefined) + return { + code: codePromise, + ready, + cancel: () => rejectOnce(new Error(`MCP OAuth authorization was cancelled for server "${this.serverId}"`)) + } + } +} + +export function createMcpOAuthProvider( + serverId: string, + server: McpServerConfig, + options: McpOAuthProviderOptions | undefined +): FileMcpOAuthProvider | undefined { + if (!options?.storageDir) return undefined + if (server.transport !== 'streamable-http' && server.transport !== 'sse') return undefined + if (!server.url || !server.oauth || server.oauth.enabled === false) return undefined + return new FileMcpOAuthProvider( + serverId, + server, + join(options.storageDir, `${safeMcpOAuthFileName(serverId)}-${hashText(server.url).slice(0, 16)}.json`), + options.openExternal, + undefined, + options.interactive ?? false, + options.encryptor + ) +} + +export async function listMcpOAuthDiagnostics( + config: McpCapabilityConfig, + options: { storageDir?: string; encryptor?: SecretEncryptor } = {} +): Promise { + const providers = createConfiguredMcpOAuthProviders(config, options) + return Promise.all(providers.map(async ({ serverId, server, provider }) => ({ + serverId, + enabled: server.oauth?.enabled !== false, + configured: Boolean(server.oauth), + transport: server.transport, + ...(server.url ? { url: server.url } : {}), + ...await provider.diagnostics() + }))) +} + +export async function clearMcpOAuthCredentials( + config: McpCapabilityConfig, + options: { storageDir?: string; serverId?: string } = {} +): Promise { + const providers = createConfiguredMcpOAuthProviders(config, options) + .filter((entry) => !options.serverId || entry.serverId === options.serverId) + await Promise.all(providers.map((entry) => entry.provider.clearCredentials())) + return { cleared: providers.map((entry) => entry.serverId) } +} + +export function createConfiguredMcpOAuthProviders( + config: McpCapabilityConfig, + options: { + storageDir?: string + openExternal?: (url: URL) => void | Promise + encryptor?: SecretEncryptor + } = {} +): Array<{ serverId: string; server: McpServerConfig; provider: FileMcpOAuthProvider }> { + return Object.entries(config.servers).flatMap(([serverId, server]) => { + const provider = createMcpOAuthProvider(serverId, server, { + storageDir: options.storageDir ?? '', + openExternal: options.openExternal, + encryptor: options.encryptor + }) + return provider ? [{ serverId, server, provider }] : [] + }) +} + +function deriveOAuthStatus(input: { + hasTokens: boolean + expired: boolean + hasPartialState: boolean + hasError: boolean +}): McpOAuthStatus { + if (input.hasTokens) return input.expired ? 'expired' : 'authorized' + if (input.hasError) return 'error' + if (input.hasPartialState) return 'partial' + return 'empty' +} + +function computeTokenExpiry(state: McpOAuthState): string | undefined { + const expiresIn = state.tokens?.expires_in + if (typeof expiresIn !== 'number' || !Number.isFinite(expiresIn)) return undefined + if (typeof state.tokensObtainedAt !== 'number' || !Number.isFinite(state.tokensObtainedAt)) return undefined + return new Date(state.tokensObtainedAt + expiresIn * 1000).toISOString() +} + +/** + * Parse the OAuth `scope` field (a space-delimited string per RFC 6749) into a + * de-duplicated list of granted scopes for the permission view. Returns an + * empty array when the provider did not echo any scope. + */ +function parseTokenScopes(scope: string | undefined): string[] { + if (typeof scope !== 'string') return [] + return [...new Set(scope.split(/\s+/).map((entry) => entry.trim()).filter(Boolean))] +} + +export function defaultMcpOAuthRedirectPort(serverId: string, url: string): number { + const digest = createHash('sha256').update(`${serverId}\n${url}`).digest() + return MCP_OAUTH_PORT_BASE + digest.readUInt16BE(0) % MCP_OAUTH_PORT_RANGE +} + +function safeMcpOAuthFileName(serverId: string): string { + return slug(serverId).slice(0, 64) || 'server' +} + +export function defaultOpenExternal(url: URL): void { + const target = url.toString() + const command = process.platform === 'win32' ? 'rundll32.exe' + : process.platform === 'darwin' ? 'open' + : 'xdg-open' + const args = process.platform === 'win32' ? ['url.dll,FileProtocolHandler', target] : [target] + const child = spawn(command, args, { + detached: true, + stdio: 'ignore', + windowsHide: true + }) + child.unref() +} diff --git a/kun/src/adapters/tool/mcp-oauth-store.test.ts b/kun/src/adapters/tool/mcp-oauth-store.test.ts new file mode 100644 index 000000000..3830bcd1f --- /dev/null +++ b/kun/src/adapters/tool/mcp-oauth-store.test.ts @@ -0,0 +1,72 @@ +import { describe, expect, it } from 'vitest' +import { mkdtemp, rm, readFile } from 'node:fs/promises' +import { tmpdir } from 'node:os' +import { join } from 'node:path' +import { randomBytes } from 'node:crypto' +import { FileMcpOAuthStore } from './mcp-oauth-store.js' +import { createAesEncryptor } from '../../security/secret-store.js' + +async function withDir(fn: (dir: string) => Promise): Promise { + const dir = await mkdtemp(join(tmpdir(), 'kun-oauth-')) + try { + return await fn(dir) + } finally { + await rm(dir, { recursive: true, force: true }) + } +} + +describe('FileMcpOAuthStore encryption', () => { + it('persists tokens encrypted at rest and decrypts on read', async () => { + await withDir(async (dir) => { + const enc = createAesEncryptor(randomBytes(32)) + const store = new FileMcpOAuthStore(join(dir, 's.json'), enc) + await store.update((s) => ({ ...s, tokens: { access_token: 'super-secret', token_type: 'bearer' } })) + const raw = await readFile(join(dir, 's.json'), 'utf8') + expect(raw).not.toContain('super-secret') + expect(raw).toContain('enc:v1:') + const read = await store.read() + expect(read.tokens?.access_token).toBe('super-secret') + }) + }) + + it('encrypts client secrets and PKCE verifier along with tokens', async () => { + await withDir(async (dir) => { + const store = new FileMcpOAuthStore(join(dir, 's.json'), createAesEncryptor(randomBytes(32))) + await store.update(() => ({ + clientInformation: { client_id: 'client', client_secret: 'client-secret' }, + codeVerifier: 'pkce-verifier', + tokens: { access_token: 'access-secret', token_type: 'bearer' } + })) + const raw = await readFile(join(dir, 's.json'), 'utf8') + expect(raw).not.toContain('client-secret') + expect(raw).not.toContain('pkce-verifier') + expect(raw).not.toContain('access-secret') + await expect(store.read()).resolves.toMatchObject({ + clientInformation: { client_secret: 'client-secret' }, + codeVerifier: 'pkce-verifier', + tokens: { access_token: 'access-secret' } + }) + }) + }) + + it('still reads legacy plaintext tokens (backward compatible)', async () => { + await withDir(async (dir) => { + const plain = new FileMcpOAuthStore(join(dir, 's.json')) + await plain.update((s) => ({ ...s, tokens: { access_token: 'legacy', token_type: 'bearer' } })) + const enc = createAesEncryptor(randomBytes(32)) + const store = new FileMcpOAuthStore(join(dir, 's.json'), enc) + const read = await store.read() + expect(read.tokens?.access_token).toBe('legacy') + }) + }) + + it('fails loudly when the credential key no longer decrypts the store', async () => { + await withDir(async (dir) => { + const store1 = new FileMcpOAuthStore(join(dir, 's.json'), createAesEncryptor(randomBytes(32))) + await store1.update((s) => ({ ...s, tokens: { access_token: 'x', token_type: 'bearer' } })) + // Different key cannot decrypt the prior blob. + const store2 = new FileMcpOAuthStore(join(dir, 's.json'), createAesEncryptor(randomBytes(32))) + await expect(store2.read()).rejects.toThrow(/could not be decrypted/) + }) + }) +}) diff --git a/kun/src/adapters/tool/mcp-oauth-store.ts b/kun/src/adapters/tool/mcp-oauth-store.ts new file mode 100644 index 000000000..3bf27e462 --- /dev/null +++ b/kun/src/adapters/tool/mcp-oauth-store.ts @@ -0,0 +1,129 @@ +import type { OAuthDiscoveryState } from '@modelcontextprotocol/sdk/client/auth.js' +import type { + OAuthClientInformationMixed, + OAuthTokens +} from '@modelcontextprotocol/sdk/shared/auth.js' +import { chmod, mkdir, readFile, rename, rm, writeFile } from 'node:fs/promises' +import { dirname } from 'node:path' +import type { SecretEncryptor } from '../../security/secret-store.js' +import { isEncryptedEnvelope } from '../../security/secret-store.js' + +/** + * Persisted OAuth credential material for a single remote MCP server. + * + * Tokens, client registration, the in-flight code verifier, and discovery + * metadata are stored on disk (mode 0600) so an authorized server keeps + * working across runtime restarts without re-prompting the user. The store + * also records lifecycle metadata used by diagnostics: + * + * - `tokensObtainedAt` lets diagnostics compute access-token expiry. + * - `lastError` / `lastErrorAt` surface the most recent failed authorization + * attempt (e.g. a loopback callback error) so the GUI can explain why a + * server is not authorized instead of showing a bare "error". + */ +export type McpOAuthState = { + clientInformation?: OAuthClientInformationMixed + tokens?: OAuthTokens + codeVerifier?: string + discoveryState?: OAuthDiscoveryState + /** Epoch milliseconds when the current tokens were saved. */ + tokensObtainedAt?: number + /** Sanitized message from the most recent failed authorization attempt. */ + lastError?: string + /** ISO timestamp of the most recent failed authorization attempt. */ + lastErrorAt?: string +} + +type EncryptedMcpOAuthState = { __enc: string } + +/** + * Atomic, single-file JSON persistence for one server's OAuth state. + * + * Writes go through a temp file + rename so a crash mid-write cannot leave a + * partially serialized credential file behind. The directory and file are + * created with restrictive permissions because they hold bearer tokens. + */ +export class FileMcpOAuthStore { + constructor( + private readonly storagePath: string, + /** Optional encryptor; when present, bearer tokens are encrypted at rest. */ + private readonly encryptor?: SecretEncryptor + ) {} + + get path(): string { + return this.storagePath + } + + async read(): Promise { + try { + const parsed = JSON.parse(await readFile(this.storagePath, 'utf8')) as unknown + if (isEncryptedMcpOAuthState(parsed)) return this.decryptEnvelope(parsed) + if (!isMcpOAuthState(parsed)) return {} + return this.decryptLegacyState(parsed) + } catch (error) { + if (isNodeErrorCode(error, 'ENOENT')) return {} + throw error + } + } + + async update(update: (state: McpOAuthState) => McpOAuthState): Promise { + const next = update(await this.read()) + await mkdir(dirname(this.storagePath), { recursive: true, mode: 0o700 }) + const temporaryPath = `${this.storagePath}.${process.pid}.${Date.now()}.tmp` + try { + await writeFile(temporaryPath, `${JSON.stringify(this.encryptState(next), null, 2)}\n`, { mode: 0o600 }) + await chmod(temporaryPath, 0o600).catch(() => undefined) + await rename(temporaryPath, this.storagePath) + await chmod(this.storagePath, 0o600).catch(() => undefined) + } finally { + await rm(temporaryPath, { force: true }).catch(() => undefined) + } + } + + async clear(): Promise { + await rm(this.storagePath, { force: true }) + } + + /** Encrypt all persisted OAuth state, including client secrets and PKCE data. */ + private encryptState(state: McpOAuthState): McpOAuthState | EncryptedMcpOAuthState { + if (!this.encryptor) return state + return { __enc: this.encryptor.encrypt(JSON.stringify(state)) } + } + + private decryptEnvelope(envelope: EncryptedMcpOAuthState): McpOAuthState { + if (!this.encryptor) throw new Error('MCP OAuth credential store is encrypted but no decryptor is available') + try { + const parsed = JSON.parse(this.encryptor.decrypt(envelope.__enc)) as unknown + if (!isMcpOAuthState(parsed)) throw new Error('decrypted OAuth state is invalid') + return parsed + } catch (error) { + throw new Error('MCP OAuth credential store could not be decrypted; restore the original key or re-authorize the connector', { cause: error }) + } + } + + /** Transparently migrate the old format that encrypted only the token field. */ + private decryptLegacyState(state: McpOAuthState): McpOAuthState { + const tokens = state.tokens as unknown as { __enc?: string } | undefined + if (!this.encryptor || !tokens || typeof tokens.__enc !== 'string') return state + if (!isEncryptedEnvelope(tokens.__enc)) return state + try { + return { ...state, tokens: JSON.parse(this.encryptor.decrypt(tokens.__enc)) as McpOAuthState['tokens'] } + } catch (error) { + throw new Error('MCP OAuth token store could not be decrypted; restore the original key or re-authorize the connector', { cause: error }) + } + } +} + +function isEncryptedMcpOAuthState(value: unknown): value is EncryptedMcpOAuthState { + return typeof value === 'object' && value !== null && !Array.isArray(value) && + Object.keys(value).length === 1 && '__enc' in value && + typeof (value as { __enc?: unknown }).__enc === 'string' +} + +export function isMcpOAuthState(value: unknown): value is McpOAuthState { + return typeof value === 'object' && value !== null && !Array.isArray(value) +} + +export function isNodeErrorCode(error: unknown, code: string): boolean { + return typeof error === 'object' && error !== null && 'code' in error && (error as { code?: unknown }).code === code +} diff --git a/kun/src/adapters/tool/mcp-stdio-environment.ts b/kun/src/adapters/tool/mcp-stdio-environment.ts new file mode 100644 index 000000000..59019f45c --- /dev/null +++ b/kun/src/adapters/tool/mcp-stdio-environment.ts @@ -0,0 +1,144 @@ +import { posix, win32 } from 'node:path' +import type { McpServerConfig } from '../../contracts/capabilities.js' + +/** + * Environment + error helpers for stdio MCP servers. + * + * Building a sane PATH for child processes (so bare `npx`/`node` commands + * resolve when Kun is launched from a GUI shell) and turning spawn failures + * into actionable messages are self-contained concerns, separated from the + * connection orchestration in `mcp-tool-provider.ts`. + */ + +export type McpStdioEnvironmentOptions = { + platform?: NodeJS.Platform + baseEnv?: NodeJS.ProcessEnv +} + +export function buildMcpStdioEnvironment( + serverEnv: Record = {}, + options: McpStdioEnvironmentOptions = {} +): Record { + const platform = options.platform ?? process.platform + const baseEnv = options.baseEnv ?? process.env + const pathKey = findPathKey(serverEnv) ?? findPathKey(baseEnv) ?? 'PATH' + const configuredPath = readEnvPath(serverEnv) + const inheritedPath = readEnvPath(baseEnv) + const pathValue = mergePathEntries( + [configuredPath ?? inheritedPath ?? '', ...commonMcpCommandPathEntries(platform, baseEnv)], + pathDelimiter(platform) + ) + return { + ...serverEnv, + ...(pathValue ? { [pathKey]: pathValue } : {}) + } +} + +export function formatMcpConnectionError(error: unknown, server: McpServerConfig): string { + const message = errorMessage(error) + if (server.transport !== 'stdio' || !isMissingExecutableError(error, message)) return message + const command = missingExecutableCommand(error) ?? server.command ?? 'configured command' + const hint = isBareCommand(command) + ? missingBareCommandHint(command) + : `Could not find MCP command "${command}". Check that the configured executable path exists.` + return `${message}. ${hint}` +} + +export function errorMessage(error: unknown): string { + return error instanceof Error ? error.message : String(error) +} + +function commonMcpCommandPathEntries( + platform: NodeJS.Platform, + env: NodeJS.ProcessEnv +): string[] { + if (platform === 'darwin') { + return [ + '/opt/homebrew/bin', + '/usr/local/bin', + '/opt/local/bin', + homePath(env, '.volta/bin'), + homePath(env, '.local/bin'), + homePath(env, '.bun/bin') + ].filter((entry): entry is string => Boolean(entry)) + } + if (platform === 'linux') { + return [ + '/home/linuxbrew/.linuxbrew/bin', + '/usr/local/bin', + '/usr/bin', + homePath(env, '.volta/bin'), + homePath(env, '.local/bin'), + homePath(env, '.bun/bin') + ].filter((entry): entry is string => Boolean(entry)) + } + if (platform === 'win32') { + return [ + env.APPDATA ? win32.join(env.APPDATA, 'npm') : '', + env.ProgramFiles ? win32.join(env.ProgramFiles, 'nodejs') : '', + env['ProgramFiles(x86)'] ? win32.join(env['ProgramFiles(x86)'], 'nodejs') : '' + ].filter((entry): entry is string => Boolean(entry)) + } + return [] +} + +function findPathKey(env: Record): string | undefined { + return Object.keys(env).find((key) => key.toLowerCase() === 'path') +} + +function readEnvPath(env: Record): string | undefined { + const key = findPathKey(env) + const value = key ? env[key] : undefined + return value && value.trim() ? value : undefined +} + +function mergePathEntries(values: string[], delimiter: string): string { + const seen = new Set() + const entries: string[] = [] + for (const value of values) { + for (const entry of value.split(delimiter)) { + const trimmed = entry.trim() + if (!trimmed) continue + const key = trimmed.toLowerCase() + if (seen.has(key)) continue + seen.add(key) + entries.push(trimmed) + } + } + return entries.join(delimiter) +} + +function pathDelimiter(platform: NodeJS.Platform): string { + return platform === 'win32' ? ';' : ':' +} + +function homePath(env: NodeJS.ProcessEnv, relativePath: string): string { + return env.HOME ? posix.join(env.HOME, relativePath) : '' +} + +function isMissingExecutableError(error: unknown, message: string): boolean { + const code = typeof error === 'object' && error !== null && 'code' in error + ? String((error as { code?: unknown }).code ?? '') + : '' + return code === 'ENOENT' || /\bspawn\s+\S+\s+ENOENT\b/i.test(message) +} + +function missingExecutableCommand(error: unknown): string | undefined { + if (!error || typeof error !== 'object') return undefined + const path = (error as { path?: unknown }).path + return typeof path === 'string' && path.trim() ? path.trim() : undefined +} + +function isBareCommand(command: string): boolean { + return Boolean(command.trim()) && !command.includes('/') && !command.includes('\\') +} + +function missingBareCommandHint(command: string): string { + if (process.platform === 'win32') { + return `Could not find "${command}" on PATH while starting the MCP server. Make sure Node/npm is installed and available to Kun, or set the MCP command to an absolute path.` + } + if (process.platform === 'darwin') { + return `Could not find "${command}" on PATH while starting the MCP server. If Kun was launched from Finder or the desktop, make sure Node/npm is installed and available to GUI apps, or set the MCP command to an absolute path such as /opt/homebrew/bin/${command}.` + } + return `Could not find "${command}" on PATH while starting the MCP server. Make sure Node/npm is installed and available to Kun, or set the MCP command to an absolute path such as /usr/local/bin/${command}.` +} diff --git a/kun/src/adapters/tool/mcp-tool-provider.ts b/kun/src/adapters/tool/mcp-tool-provider.ts index 96cc53569..4a4a67204 100644 --- a/kun/src/adapters/tool/mcp-tool-provider.ts +++ b/kun/src/adapters/tool/mcp-tool-provider.ts @@ -1,18 +1,33 @@ -import { Client } from '@modelcontextprotocol/sdk/client/index.js' -import { createHash } from 'node:crypto' -import { posix, win32 } from 'node:path' -import { SSEClientTransport } from '@modelcontextprotocol/sdk/client/sse.js' -import { StdioClientTransport } from '@modelcontextprotocol/sdk/client/stdio.js' -import { StreamableHTTPClientTransport } from '@modelcontextprotocol/sdk/client/streamableHttp.js' -import type { Transport } from '@modelcontextprotocol/sdk/shared/transport.js' -import type { - McpCapabilityConfig, - McpServerConfig -} from '../../contracts/capabilities.js' +import type { McpCapabilityConfig, McpServerConfig } from '../../contracts/capabilities.js' import { redactSecretText } from '../../config/secret-redaction.js' import type { ToolHostContext } from '../../ports/tool-host.js' import type { CapabilityToolProvider } from './capability-registry.js' import { LocalToolHost, type LocalTool } from './local-tool-host.js' +import { + catalogFingerprint, + canUseMcpServer, + isMcpServerTrusted, + isMcpServerVisible, + normalizeMcpToolName +} from './mcp-naming.js' +import { + clearMcpOAuthCredentials, + listMcpOAuthDiagnostics +} from './mcp-oauth-provider.js' +import { + authorizeMcpServerOAuth, + createSdkMcpClient, + isMcpAuthorizationRequiredError +} from './mcp-transport.js' +import { errorMessage, formatMcpConnectionError } from './mcp-stdio-environment.js' +import type { + McpClientLike, + McpOAuthAuthorizeResult, + McpOAuthClearResult, + McpOAuthDiagnostic, + McpServerDiagnostic, + McpToolDescriptor +} from './mcp-types.js' import { createMcpSearchProvider, mcpSearchDiagnostic, @@ -21,64 +36,49 @@ import { type McpSearchRuntimeDiagnostic } from './mcp-tool-search.js' -export type McpToolDescriptor = { - name: string - title?: string - description?: string - inputSchema?: Record - outputSchema?: Record - annotations?: { - title?: string - readOnlyHint?: boolean - destructiveHint?: boolean - idempotentHint?: boolean - openWorldHint?: boolean - } - execution?: unknown - icons?: unknown - _meta?: Record -} - -export type McpClientLike = { - listTools(options?: { - cursor?: string - signal?: AbortSignal - timeout?: number - }): Promise<{ tools: McpToolDescriptor[]; nextCursor?: string }> - callTool( - input: { name: string; arguments: Record }, - options?: { signal?: AbortSignal; timeout?: number } - ): Promise - close(): Promise - setLifecycleHandlers?(handlers: McpClientLifecycleHandlers): void -} - -export type McpClientLifecycleHandlers = { - onError?: (error: Error) => void - onClose?: () => void -} - -export type McpServerDiagnostic = { - id: string - enabled: boolean - transport: McpServerConfig['transport'] - trustScope: McpServerConfig['trustScope'] - available: boolean - status: 'disabled' | 'connected' | 'reconnecting' | 'error' - toolCount: number - catalogFingerprint?: string - catalogDrift?: boolean - lastConnectedAt?: string - lastDisconnectedAt?: string - lastReconnectAt?: string - nextReconnectAt?: string - reconnectAttempts?: number - lastError?: string -} +// Re-export the MCP module surface so existing consumers (and the +// `adapters/tool/index.ts` barrel) keep importing from one place even though +// the implementation now lives in focused modules: persistence, OAuth, the +// transport adapter, naming/trust, and stdio environment. +export type { + McpClientLike, + McpClientLifecycleHandlers, + McpOAuthAuthorizeResult, + McpOAuthClearResult, + McpOAuthDiagnostic, + McpOAuthStatus, + McpServerDiagnostic, + McpToolDescriptor +} from './mcp-types.js' +export { + canUseMcpServer, + isMcpServerTrusted, + isMcpServerVisible, + normalizeMcpToolName, + resolveMcpServerCwd +} from './mcp-naming.js' +export { + FileMcpOAuthProvider, + clearMcpOAuthCredentials, + createMcpOAuthProvider, + listMcpOAuthDiagnostics +} from './mcp-oauth-provider.js' +export { + authorizeMcpServerOAuth, + createSdkMcpClient, + isMcpAuthorizationRequiredError, + McpAuthorizationRequiredError +} from './mcp-transport.js' +export { + buildMcpStdioEnvironment, + formatMcpConnectionError, + type McpStdioEnvironmentOptions +} from './mcp-stdio-environment.js' export type McpToolProviderBuildResult = { providers: CapabilityToolProvider[] diagnostics: McpServerDiagnostic[] + oauth: McpOAuthDiagnostic[] search: McpSearchRuntimeDiagnostic connectedServers: number toolCount: number @@ -93,12 +93,23 @@ export type McpToolProviderBuildResult = { * failed server has reconnected or exhausted its retries (used by tests). */ startBackgroundReconnect: (register: (provider: CapabilityToolProvider) => void) => Promise + clearOAuthCredentials: (serverId?: string) => Promise + /** + * Run the interactive OAuth authorization flow for one configured remote + * server (the explicit, user-triggered entry point). Refreshes the cached + * OAuth diagnostics on completion. Startup never calls this. + */ + authorizeOAuth: (serverId: string) => Promise close: () => Promise } export type McpToolProviderOptions = { clientFactory?: (serverId: string, server: McpServerConfig) => Promise nowIso?: () => string + oauthStorageDir?: string + /** Optional encryptor so persisted OAuth tokens are encrypted at rest. */ + oauthEncryptor?: import('../../security/secret-store.js').SecretEncryptor + openExternal?: (url: URL) => void | Promise /** * Upper bound for connect + initial tool listing per server during startup. * A slow or hung server (e.g. an npx-based stdio server resolving packages) @@ -109,6 +120,12 @@ export type McpToolProviderOptions = { backgroundReconnect?: McpBackgroundReconnectOptions /** Test seam for the inter-attempt backoff; defaults to a real unref'd timer. */ delay?: (ms: number) => Promise + /** + * Test seam for the interactive authorization step. Defaults to the real + * browser-driven {@link authorizeMcpServerOAuth}. Tests inject a fake to + * exercise the authorize-then-register + reconnect path without a network. + */ + authorize?: (serverId: string, server: McpServerConfig) => Promise } export type McpBackgroundReconnectOptions = { @@ -127,11 +144,6 @@ const DEFAULT_MCP_RECONNECT_MAX_ATTEMPTS = 5 const DEFAULT_MCP_RECONNECT_BASE_DELAY_MS = 4_000 const DEFAULT_MCP_RECONNECT_MAX_DELAY_MS = 30_000 -export type McpStdioEnvironmentOptions = { - platform?: NodeJS.Platform - baseEnv?: NodeJS.ProcessEnv -} - type McpConnectionState = { serverId: string server: McpServerConfig @@ -141,14 +153,18 @@ type McpConnectionState = { catalogFingerprint?: string catalogDrift?: boolean lastConnectedAt?: string - lastDisconnectedAt?: string - lastReconnectAt?: string - nextReconnectAt?: string + lastError?: string + // Reconnect state machine (#642/#639), ported from upstream so a dropped + // transport flips the live diagnostic to `reconnecting`/`error` and a single + // shared reconnect recovers concurrent callers. + status: 'connected' | 'reconnecting' | 'error' reconnectAttempts: number reconnectBackoffMs: number reconnectPromise?: Promise - lastError?: string - status: 'connected' | 'reconnecting' | 'error' + lastDisconnectedAt?: string + lastReconnectAt?: string + nextReconnectAt?: string + /** Live diagnostic object — the SAME reference stored in the diagnostics array. */ diagnostic?: McpServerDiagnostic intentionallyClosing?: boolean } @@ -164,11 +180,17 @@ export async function buildMcpToolProviders( const catalogState: McpSearchCatalogState = { records: [] } const mcp = config const nowIso = options.nowIso ?? (() => new Date().toISOString()) - const clientFactory = options.clientFactory ?? createSdkMcpClient + const clientFactory = options.clientFactory ?? ((serverId, server) => + createSdkMcpClient(serverId, server, { + storageDir: options.oauthStorageDir, + openExternal: options.openExternal, + ...(options.oauthEncryptor ? { encryptor: options.oauthEncryptor } : {}) + })) if (!mcp?.enabled) { return { providers, diagnostics, + oauth: [], search: mcpSearchDiagnostic({ config: config?.search ?? { enabled: false, @@ -187,6 +209,8 @@ export async function buildMcpToolProviders( connectedServers: 0, toolCount: 0, startBackgroundReconnect: async () => undefined, + clearOAuthCredentials: async () => ({ cleared: [] }), + authorizeOAuth: async (serverId) => ({ serverId, status: 'disabled', authorized: false }), close: async () => undefined } } @@ -218,9 +242,9 @@ export async function buildMcpToolProviders( client, clientFactory, nowIso, + status: 'connected', reconnectAttempts: 0, reconnectBackoffMs: DEFAULT_MCP_RECONNECT_BASE_DELAY_MS, - status: 'connected', lastConnectedAt: nowIso() } attachMcpClientLifecycle(state) @@ -242,12 +266,13 @@ export async function buildMcpToolProviders( continue } if (outcome.status === 'error') { + const authRequired = isMcpAuthorizationRequiredError(outcome.error) diagnostics.push( serverDiagnostic( { serverId: outcome.serverId, server: outcome.server }, - 'error', + authRequired ? 'authorization_required' : 'error', 0, - formatMcpConnectionError(outcome.error, outcome.server) + startupConnectionError(outcome.error, outcome.server) ) ) continue @@ -268,6 +293,10 @@ export async function buildMcpToolProviders( const connectedServers = diagnostics.filter((diagnostic) => diagnostic.status === 'connected').length const toolCount = catalogState.records.length + const oauthDiagnostics = await listMcpOAuthDiagnostics(mcp, { + storageDir: options.oauthStorageDir, + encryptor: options.oauthEncryptor + }) catalogState.lastRefreshedAt = nowIso() catalogState.catalogFingerprint = catalogFingerprint(catalogState.records.map((record) => record.toolId)) const searchActive = shouldUseMcpSearch(mcp.search, toolCount) && connectedServers > 0 @@ -301,15 +330,113 @@ export async function buildMcpToolProviders( } const advertisedToolCount = providers.reduce((total, provider) => total + provider.tools.length, 0) + // Servers that need OAuth authorization are NOT retried by the background + // reconnect loop — retrying just burns attempts and would re-hit a 401. They + // wait in `authorization_required` until the user authorizes, after which + // authorizeOAuth() performs a single live connect + register. const failedServers = outcomes.flatMap((outcome) => - outcome.status === 'error' ? [{ serverId: outcome.serverId, server: outcome.server }] : [] + outcome.status === 'error' && !isMcpAuthorizationRequiredError(outcome.error) + ? [{ serverId: outcome.serverId, server: outcome.server }] + : [] ) let reconnectAborted = false let reconnectStarted = false + /** Captured from startBackgroundReconnect so authorizeOAuth can register live. */ + let liveRegister: ((provider: CapabilityToolProvider) => void) | null = null + /** Per-serverId authorization single-flight: concurrent clicks share one run. */ + const authorizeInFlight = new Map>() + + const refreshOAuthDiagnostics = async (): Promise => { + const nextDiagnostics = await listMcpOAuthDiagnostics(mcp, { + storageDir: options.oauthStorageDir, + encryptor: options.oauthEncryptor + }) + oauthDiagnostics.splice(0, oauthDiagnostics.length, ...nextDiagnostics) + } + + /** + * Connect a server live (using the real/injected client factory), list its + * tools, register the provider, and flip its diagnostic to `connected` — no + * runtime restart required after a successful authorization. + */ + const connectAndRegisterServer = async (serverId: string, server: McpServerConfig): Promise => { + const client = await clientFactory(serverId, server) + const state: McpConnectionState = { + serverId, + server, + client, + clientFactory, + nowIso, + status: 'connected', + reconnectAttempts: 0, + reconnectBackoffMs: DEFAULT_MCP_RECONNECT_BASE_DELAY_MS, + lastConnectedAt: nowIso() + } + attachMcpClientLifecycle(state) + let listed: McpToolDescriptor[] + try { + listed = await refreshMcpConnectionCatalog(state) + } catch (error) { + await client.close().catch(() => undefined) + throw error + } + connected.push(state) + catalogState.records.push(...listed.map((tool) => createMcpSearchCatalogRecord(state, tool))) + catalogState.catalogFingerprint = catalogFingerprint(catalogState.records.map((record) => record.toolId)) + catalogState.lastRefreshedAt = nowIso() + const tools = listed.map((tool) => createMcpLocalTool(state, tool)) + if (!searchActive && liveRegister) { + try { + liveRegister({ id: `mcp:${serverId}`, kind: 'mcp', enabled: true, available: true, tools }) + } catch { + // Registry collision must not break the authorize flow; the diagnostic + // still flips to connected below. + } + } + const diagnostic = syncMcpDiagnostic(state, 'connected', tools.length) + const index = diagnostics.findIndex((entry) => entry.id === serverId) + if (index >= 0) diagnostics[index] = diagnostic + else diagnostics.push(diagnostic) + } + + const authorizeOAuth = (serverId: string): Promise => { + const inflight = authorizeInFlight.get(serverId) + if (inflight) return inflight + const run = (async (): Promise => { + const server = mcp.servers[serverId] + if (!server || !options.oauthStorageDir) { + return { serverId, status: 'disabled', authorized: false } + } + const authorize = options.authorize ?? + ((id: string, srv: McpServerConfig) => authorizeMcpServerOAuth(id, srv, { + storageDir: options.oauthStorageDir as string, + openExternal: options.openExternal, + encryptor: options.oauthEncryptor + })) + const result = await authorize(serverId, server) + await refreshOAuthDiagnostics() + // On success, connect + register immediately so tools are live without a + // runtime restart. Skip if the server is already connected. + if (result.authorized && !connected.some((state) => state.serverId === serverId)) { + try { + await connectAndRegisterServer(serverId, server) + } catch { + // Leave the server in its prior diagnostic state; the user can retry. + } + } + return result + })() + authorizeInFlight.set(serverId, run) + run.finally(() => { + if (authorizeInFlight.get(serverId) === run) authorizeInFlight.delete(serverId) + }).catch(() => undefined) + return run + } return { providers, diagnostics, + oauth: oauthDiagnostics, search: mcpSearchDiagnostic({ config: mcp.search, active: searchActive, @@ -320,6 +447,7 @@ export async function buildMcpToolProviders( connectedServers, toolCount, startBackgroundReconnect: (register) => { + liveRegister = register if (reconnectStarted) return Promise.resolve() reconnectStarted = true if (failedServers.length === 0) return Promise.resolve() @@ -338,13 +466,35 @@ export async function buildMcpToolProviders( options: options.backgroundReconnect }) }, + clearOAuthCredentials: async (serverId) => { + const result = await clearMcpOAuthCredentials(mcp, { + storageDir: options.oauthStorageDir, + serverId + }) + await refreshOAuthDiagnostics() + return result + }, + authorizeOAuth, close: async () => { reconnectAborted = true - await Promise.all(connected.map((state) => closeMcpClient(state))) + await Promise.all(connected.map((state) => state.client.close().catch(() => undefined))) } } } +/** + * Turn a startup connect failure into an actionable diagnostic message. + * Authorization-required failures (a remote OAuth server with no usable token) + * are expected during startup — the connect is non-interactive — so they get a + * "use Authorize" hint instead of a raw transport error. + */ +function startupConnectionError(error: unknown, server: McpServerConfig): string { + if (isMcpAuthorizationRequiredError(error)) { + return 'OAuth authorization required. Use the connector\'s Authorize action to sign in; the runtime will not prompt automatically during startup.' + } + return formatMcpConnectionError(error, server) +} + type FailedMcpServer = { serverId: string; server: McpServerConfig } type McpBackgroundReconnectParams = { @@ -389,7 +539,6 @@ async function reconnectFailedMcpServer( ): Promise { for (let attempt = 1; attempt <= maxAttempts; attempt += 1) { if (params.isAborted()) return - updateFailedServerDiagnostic(params.diagnostics, failed, 'reconnecting', attempt) await params.delay(Math.min(maxDelayMs, baseDelayMs * 2 ** (attempt - 1))) if (params.isAborted()) return try { @@ -400,27 +549,21 @@ async function reconnectFailedMcpServer( client, clientFactory: params.clientFactory, nowIso: params.nowIso, + status: 'connected', reconnectAttempts: 0, reconnectBackoffMs: DEFAULT_MCP_RECONNECT_BASE_DELAY_MS, - status: 'connected', lastConnectedAt: params.nowIso() } attachMcpClientLifecycle(state) const listed = await refreshMcpConnectionCatalog(state) if (params.isAborted()) { - await closeMcpClient(state) + await client.close().catch(() => undefined) return } registerLateMcpConnection(params, state, listed) return - } catch (error) { - updateFailedServerDiagnostic( - params.diagnostics, - failed, - 'error', - attempt, - formatMcpConnectionError(error, failed.server) - ) + } catch { + // Leave the diagnostic as "error" and try again until attempts run out. } } } @@ -456,27 +599,6 @@ function registerLateMcpConnection( else params.diagnostics.push(diagnostic) } -function updateFailedServerDiagnostic( - diagnostics: McpServerDiagnostic[], - failed: FailedMcpServer, - status: Extract, - attempt: number, - lastError?: string -): void { - const index = diagnostics.findIndex((entry) => entry.id === failed.serverId) - const previous = index >= 0 ? diagnostics[index] : undefined - const next: McpServerDiagnostic = { - ...(previous ?? serverDiagnostic({ serverId: failed.serverId, server: failed.server }, 'error', 0)), - status, - available: false, - reconnectAttempts: attempt, - lastReconnectAt: new Date().toISOString(), - ...(lastError ? { lastError: redactSecretText(lastError) } : {}) - } - if (index >= 0) diagnostics[index] = next - else diagnostics.push(next) -} - function defaultMcpReconnectDelay(ms: number): Promise { return new Promise((resolve) => { const timer = setTimeout(resolve, ms) @@ -486,105 +608,6 @@ function defaultMcpReconnectDelay(ms: number): Promise { }) } -export function normalizeMcpToolName(serverId: string, toolName: string): string { - return `mcp_${slug(serverId)}_${slug(toolName)}` -} - -export function isMcpServerTrusted(server: McpServerConfig, workspace: string): boolean { - if (server.trustScope === 'user') return true - return workspaceMatchesRoots(workspace, server.trustedWorkspaceRoots) -} - -export function isMcpServerVisible(server: McpServerConfig, workspace: string): boolean { - if (server.workspaceRoots.length === 0) return true - return workspaceMatchesRoots(workspace, server.workspaceRoots) -} - -export function canUseMcpServer(server: McpServerConfig, workspace: string): boolean { - return isMcpServerVisible(server, workspace) && isMcpServerTrusted(server, workspace) -} - -function workspaceMatchesRoots(workspace: string, roots: readonly string[]): boolean { - const normalizedWorkspace = normalizePathForTrust(workspace) - return roots.some((root) => { - const normalizedRoot = normalizePathForTrust(root) - return normalizedWorkspace === normalizedRoot || normalizedWorkspace.startsWith(`${normalizedRoot}/`) - }) -} - -async function createSdkMcpClient(serverId: string, server: McpServerConfig): Promise { - const client = new Client({ name: `kun-${serverId}`, version: '0.1.0' }) - // Observe transport-level failures explicitly. The SDK routes a dropped SSE - // stream and exhausted background reconnects to `onerror`; with no handler - // they are silently swallowed, which hides real outages from the logs and (on - // some SDK/runtime versions) lets the rejection escape as unhandled. Handling - // it here keeps a streamable-http disconnect from destabilizing the runtime - // (#639) — the per-call reconnect in callMcpToolWithReconnect still recovers - // the connection on the next tool use. - client.onerror = (error) => { - process.stderr.write(`kun mcp[${serverId}]: transport error: ${redactSecretText(errorMessage(error))}\n`) - } - const transport = createTransport(server) - await client.connect(transport, { timeout: server.timeoutMs }) - return { - listTools: (options) => { - const params = options?.cursor ? { cursor: options.cursor } : undefined - return client.listTools(params, { - signal: options?.signal, - timeout: options?.timeout - }) - }, - callTool: (input, options) => client.callTool(input, undefined, options), - close: () => client.close(), - setLifecycleHandlers: (handlers) => { - client.onerror = handlers.onError - client.onclose = handlers.onClose - } - } -} - -function createTransport(server: McpServerConfig): Transport { - switch (server.transport) { - case 'stdio': { - const cwd = resolveMcpServerCwd(server) - return new StdioClientTransport({ - command: server.command ?? '', - args: server.args, - env: buildMcpStdioEnvironment(server.env), - ...(cwd ? { cwd } : {}), - stderr: 'pipe' - }) - } - case 'streamable-http': - return new StreamableHTTPClientTransport(new URL(server.url ?? ''), { - requestInit: { headers: server.headers } - }) - case 'sse': - return new SSEClientTransport(new URL(server.url ?? ''), { - requestInit: { headers: server.headers }, - eventSourceInit: { fetch: fetchWithHeaders(server.headers) } - }) - } -} - -export function resolveMcpServerCwd(server: McpServerConfig): string | undefined { - if (server.transport !== 'stdio') return undefined - const configured = server.cwd?.trim() - if (configured) return configured - if (server.trustScope !== 'workspace') return undefined - return server.trustedWorkspaceRoots.map((root) => root.trim()).find(Boolean) -} - -function fetchWithHeaders(headers: Record): typeof fetch { - return (input, init) => { - const mergedHeaders = new Headers(init?.headers) - for (const [key, value] of Object.entries(headers)) { - mergedHeaders.set(key, value) - } - return fetch(input, { ...init, headers: mergedHeaders }) - } -} - function createMcpLocalTool( state: McpConnectionState, descriptor: McpToolDescriptor @@ -608,21 +631,12 @@ function createMcpLocalTool( isError: true } } - if (state.status === 'error' && !canAttemptMcpReconnect(state)) { - return { - output: { - error: mcpReconnectCooldownMessage(state), - serverId: state.serverId, - status: state.status, - nextReconnectAt: state.nextReconnectAt - }, - isError: true - } - } const result = await callMcpToolWithReconnect( state, { name: descriptor.name, arguments: args }, - context.abortSignal + context.abortSignal, + state.server.timeoutMs, + isMcpReplaySafe(descriptor.annotations) ) return { output: { @@ -657,7 +671,7 @@ function createMcpSearchCatalogRecord( server: state.server, client: { callTool: (input, options) => - callMcpToolWithReconnect(state, input, options?.signal, options?.timeout) + callMcpToolWithReconnect(state, input, options?.signal, options?.timeout, isMcpReplaySafe(descriptor.annotations)) }, descriptor, normalizedName: normalizeMcpToolName(state.serverId, descriptor.name), @@ -679,10 +693,22 @@ async function callMcpToolWithReconnect( state: McpConnectionState, input: { name: string; arguments: Record }, signal: AbortSignal | undefined, - timeout = state.server.timeoutMs + timeout = state.server.timeoutMs, + /** + * Whether this specific tool is safe to REPLAY after a mid-call transport + * drop. Only tools explicitly marked read-only or idempotent are safe; a + * non-idempotent tool may have already executed on the server before the + * transport failed, so replaying it could duplicate its side effects. + */ + replaySafe = false ): Promise { + // Track whether the request actually reached `callTool`. A failure while + // (re)connecting BEFORE the request was sent means the tool definitely did + // not run, so retrying it on the fresh connection is always safe. + let sentToServer = false try { await ensureMcpConnectionForCall(state, signal) + sentToServer = true return await state.client.callTool(input, { signal, timeout }) } catch (error) { if (signal?.aborted) throw error @@ -696,8 +722,52 @@ async function callMcpToolWithReconnect( throw error } markMcpConnectionError(state, error) - const client = await reconnectMcpConnection(state, signal) - return client.callTool(input, { signal, timeout }) + if (!sentToServer || replaySafe) { + // Either the request never left (safe) or the tool is read-only/idempotent + // (safe to repeat) — replay it on the reconnected client. + const client = await reconnectMcpConnection(state, signal) + return client.callTool(input, { signal, timeout }) + } + // A non-idempotent tool dropped mid-flight: it may already have run on the + // server. Preserve that primary fact even when reconnect also fails: future + // calls can reconnect in the background, but this call must always surface + // status-unknown rather than a secondary transport error. + void reconnectMcpConnection(state, undefined).catch(() => undefined) + throw new McpToolStatusUnknownError(state.serverId, input.name, error) + } +} + +/** + * MCP annotations that make a tool safe to replay after a mid-call transport + * drop. Per the MCP spec the hints default to the UNSAFE side (readOnlyHint + * false, idempotentHint false, destructiveHint effectively true), so we only + * treat a tool as replay-safe when it is EXPLICITLY read-only or idempotent. + * Absent annotations → not safe → no auto-replay. + */ +function isMcpReplaySafe(annotations: McpToolDescriptor['annotations']): boolean { + if (!annotations) return false + if (annotations.readOnlyHint === true && annotations.destructiveHint !== true) return true + if (annotations.idempotentHint === true && annotations.destructiveHint !== true) return true + return false +} + +/** + * Thrown when a non-idempotent MCP tool's transport dropped mid-call, so its + * server-side outcome is unknown and it was NOT auto-replayed. + */ +export class McpToolStatusUnknownError extends Error { + readonly statusUnknown = true + constructor( + readonly serverId: string, + readonly toolName: string, + readonly causeError: unknown + ) { + super( + `MCP tool "${toolName}" on server "${serverId}" lost its connection mid-call; ` + + 'its result is unknown and it was not retried automatically because it is not marked read-only or idempotent. ' + + 'Verify whether it took effect before re-running it.' + ) + this.name = 'McpToolStatusUnknownError' } } @@ -906,6 +976,9 @@ function syncMcpDiagnostic( ...(state.reconnectAttempts > 0 ? { reconnectAttempts: state.reconnectAttempts } : {}), ...(state.lastError ? { lastError: redactSecretText(state.lastError) } : {}) } + // The diagnostics array stores this exact object reference; mutate it in + // place so live status changes (reconnecting/error/connected) are visible to + // anyone holding the array without re-indexing. if (!state.diagnostic) { state.diagnostic = diagnostic return diagnostic @@ -916,173 +989,3 @@ function syncMcpDiagnostic( Object.assign(state.diagnostic, diagnostic) return state.diagnostic } - -function catalogFingerprint(values: readonly string[]): string { - return createHash('sha256') - .update(JSON.stringify([...values].sort())) - .digest('hex') - .slice(0, 16) -} - -function slug(value: string): string { - let out = '' - for (const char of value.trim().toLowerCase()) { - if (isSlugChar(char)) { - out += char - } else if (out && out[out.length - 1] !== '_') { - out += '_' - } - } - return trimBoundaryUnderscores(out) || 'tool' -} - -function normalizePathForTrust(value: string): string { - return trimTrailingSlashes(value.replaceAll('\\', '/')) -} - -function isSlugChar(char: string): boolean { - const code = char.charCodeAt(0) - return char === '_' || (code >= 48 && code <= 57) || (code >= 97 && code <= 122) -} - -function trimBoundaryUnderscores(value: string): string { - let start = 0 - let end = value.length - while (start < end && value[start] === '_') start += 1 - while (end > start && value[end - 1] === '_') end -= 1 - return value.slice(start, end) -} - -function trimTrailingSlashes(value: string): string { - let end = value.length - while (end > 0 && value.charCodeAt(end - 1) === 47) end -= 1 - return end === value.length ? value : value.slice(0, end) -} - -function errorMessage(error: unknown): string { - return error instanceof Error ? error.message : String(error) -} - -export function buildMcpStdioEnvironment( - serverEnv: Record = {}, - options: McpStdioEnvironmentOptions = {} -): Record { - const platform = options.platform ?? process.platform - const baseEnv = options.baseEnv ?? process.env - const pathKey = findPathKey(serverEnv) ?? findPathKey(baseEnv) ?? 'PATH' - const configuredPath = readEnvPath(serverEnv) - const inheritedPath = readEnvPath(baseEnv) - const pathValue = mergePathEntries( - [configuredPath ?? inheritedPath ?? '', ...commonMcpCommandPathEntries(platform, baseEnv)], - pathDelimiter(platform) - ) - return { - ...serverEnv, - ...(pathValue ? { [pathKey]: pathValue } : {}) - } -} - -export function formatMcpConnectionError(error: unknown, server: McpServerConfig): string { - const message = errorMessage(error) - if (server.transport !== 'stdio' || !isMissingExecutableError(error, message)) return message - const command = missingExecutableCommand(error) ?? server.command ?? 'configured command' - const hint = isBareCommand(command) - ? missingBareCommandHint(command) - : `Could not find MCP command "${command}". Check that the configured executable path exists.` - return `${message}. ${hint}` -} - -function commonMcpCommandPathEntries( - platform: NodeJS.Platform, - env: NodeJS.ProcessEnv -): string[] { - if (platform === 'darwin') { - return [ - '/opt/homebrew/bin', - '/usr/local/bin', - '/opt/local/bin', - homePath(env, '.volta/bin'), - homePath(env, '.local/bin'), - homePath(env, '.bun/bin') - ].filter((entry): entry is string => Boolean(entry)) - } - if (platform === 'linux') { - return [ - '/home/linuxbrew/.linuxbrew/bin', - '/usr/local/bin', - '/usr/bin', - homePath(env, '.volta/bin'), - homePath(env, '.local/bin'), - homePath(env, '.bun/bin') - ].filter((entry): entry is string => Boolean(entry)) - } - if (platform === 'win32') { - return [ - env.APPDATA ? win32.join(env.APPDATA, 'npm') : '', - env.ProgramFiles ? win32.join(env.ProgramFiles, 'nodejs') : '', - env['ProgramFiles(x86)'] ? win32.join(env['ProgramFiles(x86)'], 'nodejs') : '' - ].filter((entry): entry is string => Boolean(entry)) - } - return [] -} - -function findPathKey(env: Record): string | undefined { - return Object.keys(env).find((key) => key.toLowerCase() === 'path') -} - -function readEnvPath(env: Record): string | undefined { - const key = findPathKey(env) - const value = key ? env[key] : undefined - return value && value.trim() ? value : undefined -} - -function mergePathEntries(values: string[], delimiter: string): string { - const seen = new Set() - const entries: string[] = [] - for (const value of values) { - for (const entry of value.split(delimiter)) { - const trimmed = entry.trim() - if (!trimmed) continue - const key = trimmed.toLowerCase() - if (seen.has(key)) continue - seen.add(key) - entries.push(trimmed) - } - } - return entries.join(delimiter) -} - -function pathDelimiter(platform: NodeJS.Platform): string { - return platform === 'win32' ? ';' : ':' -} - -function homePath(env: NodeJS.ProcessEnv, relativePath: string): string { - return env.HOME ? posix.join(env.HOME, relativePath) : '' -} - -function isMissingExecutableError(error: unknown, message: string): boolean { - const code = typeof error === 'object' && error !== null && 'code' in error - ? String((error as { code?: unknown }).code ?? '') - : '' - return code === 'ENOENT' || /\bspawn\s+\S+\s+ENOENT\b/i.test(message) -} - -function missingExecutableCommand(error: unknown): string | undefined { - if (!error || typeof error !== 'object') return undefined - const path = (error as { path?: unknown }).path - return typeof path === 'string' && path.trim() ? path.trim() : undefined -} - -function isBareCommand(command: string): boolean { - return Boolean(command.trim()) && !command.includes('/') && !command.includes('\\') -} - -function missingBareCommandHint(command: string): string { - if (process.platform === 'win32') { - return `Could not find "${command}" on PATH while starting the MCP server. Make sure Node/npm is installed and available to Kun, or set the MCP command to an absolute path.` - } - if (process.platform === 'darwin') { - return `Could not find "${command}" on PATH while starting the MCP server. If Kun was launched from Finder or the desktop, make sure Node/npm is installed and available to GUI apps, or set the MCP command to an absolute path such as /opt/homebrew/bin/${command}.` - } - return `Could not find "${command}" on PATH while starting the MCP server. Make sure Node/npm is installed and available to Kun, or set the MCP command to an absolute path such as /usr/local/bin/${command}.` -} diff --git a/kun/src/adapters/tool/mcp-transport.ts b/kun/src/adapters/tool/mcp-transport.ts new file mode 100644 index 000000000..628c22d88 --- /dev/null +++ b/kun/src/adapters/tool/mcp-transport.ts @@ -0,0 +1,197 @@ +import { Client } from '@modelcontextprotocol/sdk/client/index.js' +import { UnauthorizedError, type OAuthClientProvider } from '@modelcontextprotocol/sdk/client/auth.js' +import { SSEClientTransport } from '@modelcontextprotocol/sdk/client/sse.js' +import { StdioClientTransport } from '@modelcontextprotocol/sdk/client/stdio.js' +import { StreamableHTTPClientTransport } from '@modelcontextprotocol/sdk/client/streamableHttp.js' +import type { Transport } from '@modelcontextprotocol/sdk/shared/transport.js' +import type { McpServerConfig } from '../../contracts/capabilities.js' +import { resolveMcpServerCwd } from './mcp-naming.js' +import { + FileMcpOAuthProvider, + createMcpOAuthProvider, + defaultOpenExternal +} from './mcp-oauth-provider.js' +import { buildMcpStdioEnvironment, errorMessage } from './mcp-stdio-environment.js' +import { redactSecretText } from '../../config/secret-redaction.js' +import { + McpAuthorizationRequiredError, + isMcpAuthorizationRequiredError, + type McpClientLike, + type McpOAuthAuthorizeResult +} from './mcp-types.js' + +type OAuthTransport = Transport & { + finishAuth?: (authorizationCode: string) => Promise +} + +export type SdkMcpClientOptions = { + storageDir?: string + openExternal?: (url: URL) => void | Promise + /** + * Whether an unauthorized remote server may launch the interactive browser + * authorization flow. Defaults to `false` so the startup connect pass never + * blocks the runtime on a user completing an OAuth handshake; an explicit, + * user-triggered authorize call passes `true`. The flag is forwarded to the + * OAuth provider so the browser/callback is refused at the source, not just + * after the connect throws. + */ + interactive?: boolean + /** Optional encryptor; when present, persisted OAuth tokens are encrypted at rest. */ + encryptor?: import('../../security/secret-store.js').SecretEncryptor +} + +export { McpAuthorizationRequiredError, isMcpAuthorizationRequiredError } + +export async function createSdkMcpClient( + serverId: string, + server: McpServerConfig, + options: SdkMcpClientOptions = {} +): Promise { + const client = new Client({ name: `kun-${serverId}`, version: '0.1.0' }) + // Observe transport-level failures explicitly (#639). The SDK routes a + // dropped SSE stream / exhausted reconnect to `onerror`; with no handler it + // is silently swallowed and can escape as an unhandled rejection that + // destabilizes the runtime. A default logging handler keeps a streamable-http + // disconnect from taking down the runtime; per-call reconnect still recovers + // the connection on the next tool use. `setLifecycleHandlers` lets the + // orchestration override this to drive the reconnect state machine. + ;(client as { onerror?: (error: Error) => void }).onerror = (error) => { + process.stderr.write(`kun mcp[${serverId}]: transport error: ${redactSecretText(errorMessage(error))}\n`) + } + const authProvider = createMcpOAuthProvider(serverId, server, { + storageDir: options.storageDir ?? '', + openExternal: options.openExternal, + interactive: options.interactive ?? false, + ...(options.encryptor ? { encryptor: options.encryptor } : {}) + }) + let transport = createTransport(server, authProvider) + try { + await client.connect(transport, { timeout: server.timeoutMs }) + } catch (error) { + // The non-interactive provider throws this synchronously from + // redirectToAuthorization (before any browser/callback), so it surfaces + // here as the connect rejection. Propagate it untouched. + if (isMcpAuthorizationRequiredError(error)) { + await client.close().catch(() => undefined) + throw error + } + if (!(error instanceof UnauthorizedError) || !authProvider || typeof transport.finishAuth !== 'function') { + await client.close().catch(() => undefined) + throw error + } + // Defensive: an UnauthorizedError on a non-interactive connect means the + // SDK reached auth() without the provider throwing — still never open a + // browser; report "needs authorization". + if (!options.interactive) { + await client.close().catch(() => undefined) + throw new McpAuthorizationRequiredError(serverId) + } + try { + const authorizationCode = await authProvider.waitForAuthorizationCode() + await transport.finishAuth(authorizationCode) + } catch (authError) { + await authProvider.recordAuthorizationError(errorMessage(authError)).catch(() => undefined) + await client.close().catch(() => undefined) + throw authError + } + transport = createTransport(server, authProvider) + await client.connect(transport, { timeout: server.timeoutMs }) + } + return { + listTools: (listOptions) => { + const params = listOptions?.cursor ? { cursor: listOptions.cursor } : undefined + return client.listTools(params, { + signal: listOptions?.signal, + timeout: listOptions?.timeout + }) + }, + callTool: (input, callOptions) => client.callTool(input, undefined, callOptions), + close: () => client.close(), + setLifecycleHandlers: (handlers) => { + ;(client as { onerror?: (error: Error) => void }).onerror = handlers.onError + ;(client as { onclose?: () => void }).onclose = handlers.onClose + } + } +} + +/** + * Run the interactive OAuth authorization flow for one remote MCP server and + * report the resulting credential status. This is the explicit, user-triggered + * entry point: it opens the browser, waits for the loopback callback, persists + * the tokens, then tears the probe connection down. It never participates in + * the runtime startup path. + */ +export async function authorizeMcpServerOAuth( + serverId: string, + server: McpServerConfig, + options: { + storageDir: string + openExternal?: (url: URL) => void | Promise + encryptor?: import('../../security/secret-store.js').SecretEncryptor + } +): Promise { + const provider = createMcpOAuthProvider(serverId, server, { + storageDir: options.storageDir, + openExternal: options.openExternal ?? defaultOpenExternal, + interactive: true, + encryptor: options.encryptor + }) + if (!provider) { + return { serverId, status: 'disabled', authorized: false } + } + let client: McpClientLike | undefined + try { + client = await createSdkMcpClient(serverId, server, { + storageDir: options.storageDir, + openExternal: options.openExternal ?? defaultOpenExternal, + interactive: true, + encryptor: options.encryptor + }) + } catch (error) { + await provider.recordAuthorizationError(errorMessage(error)).catch(() => undefined) + } finally { + await client?.close().catch(() => undefined) + } + const diagnostics = await provider.diagnostics() + return { serverId, status: diagnostics.status, authorized: diagnostics.status === 'authorized' } +} + +export function createTransport(server: McpServerConfig, authProvider?: OAuthClientProvider): OAuthTransport { + switch (server.transport) { + case 'stdio': { + const cwd = resolveMcpServerCwd(server) + return new StdioClientTransport({ + command: server.command ?? '', + args: server.args, + env: buildMcpStdioEnvironment(server.env), + ...(cwd ? { cwd } : {}), + stderr: 'pipe' + }) + } + case 'streamable-http': + return new StreamableHTTPClientTransport(new URL(server.url ?? ''), { + ...(authProvider ? { authProvider } : {}), + requestInit: { headers: server.headers } + }) + case 'sse': + return new SSEClientTransport(new URL(server.url ?? ''), { + ...(authProvider ? { authProvider } : {}), + requestInit: { headers: server.headers }, + eventSourceInit: { fetch: fetchWithHeaders(server.headers) } + }) + } +} + +export function fetchWithHeaders(headers: Record): typeof fetch { + return (input, init) => { + const mergedHeaders = new Headers(init?.headers) + for (const [key, value] of Object.entries(headers)) { + mergedHeaders.set(key, value) + } + return fetch(input, { ...init, headers: mergedHeaders }) + } +} + +// Re-export so callers that build providers directly (e.g. an authorize flow) +// can stay on this module's surface. +export { FileMcpOAuthProvider } diff --git a/kun/src/adapters/tool/mcp-types.ts b/kun/src/adapters/tool/mcp-types.ts new file mode 100644 index 000000000..013a81d74 --- /dev/null +++ b/kun/src/adapters/tool/mcp-types.ts @@ -0,0 +1,97 @@ +import type { McpServerConfig } from '../../contracts/capabilities.js' + +export type McpToolDescriptor = { + name: string + title?: string + description?: string + inputSchema?: Record + outputSchema?: Record + annotations?: { + title?: string + readOnlyHint?: boolean + destructiveHint?: boolean + idempotentHint?: boolean + openWorldHint?: boolean + } + execution?: unknown + icons?: unknown + _meta?: Record +} + +export type McpClientLifecycleHandlers = { + onError?: (error: Error) => void + onClose?: () => void +} + +export type McpClientLike = { + listTools(options?: { + cursor?: string + signal?: AbortSignal + timeout?: number + }): Promise<{ tools: McpToolDescriptor[]; nextCursor?: string }> + callTool( + input: { name: string; arguments: Record }, + options?: { signal?: AbortSignal; timeout?: number } + ): Promise + close(): Promise + setLifecycleHandlers?(handlers: McpClientLifecycleHandlers): void +} + +export type McpServerDiagnostic = { + id: string + enabled: boolean + transport: McpServerConfig['transport'] + trustScope: McpServerConfig['trustScope'] + available: boolean + status: 'disabled' | 'connected' | 'reconnecting' | 'error' | 'authorization_required' + toolCount: number + catalogFingerprint?: string + catalogDrift?: boolean + lastConnectedAt?: string + lastDisconnectedAt?: string + lastReconnectAt?: string + nextReconnectAt?: string + reconnectAttempts?: number + lastError?: string +} + +export type McpOAuthStatus = 'disabled' | 'empty' | 'partial' | 'authorized' | 'expired' | 'error' + +export type McpOAuthDiagnostic = { + serverId: string + enabled: boolean + configured: boolean + transport: McpServerConfig['transport'] + url?: string + status: McpOAuthStatus + hasClientInformation: boolean + hasTokens: boolean + hasRefreshToken: boolean + hasCodeVerifier: boolean + hasDiscoveryState: boolean + grantedScopes?: string[] + expiresAt?: string + lastError?: string + lastErrorAt?: string +} + +export type McpOAuthClearResult = { + cleared: string[] +} + +export type McpOAuthAuthorizeResult = { + serverId: string + status: McpOAuthStatus + authorized: boolean +} + +export class McpAuthorizationRequiredError extends Error { + constructor(public readonly serverId: string) { + super(`MCP server "${serverId}" requires OAuth authorization`) + this.name = 'McpAuthorizationRequiredError' + } +} + +export function isMcpAuthorizationRequiredError(error: unknown): error is McpAuthorizationRequiredError { + return error instanceof McpAuthorizationRequiredError +} diff --git a/kun/src/adapters/tool/memory-tool-provider.ts b/kun/src/adapters/tool/memory-tool-provider.ts index f7448ddce..45c8c85a4 100644 --- a/kun/src/adapters/tool/memory-tool-provider.ts +++ b/kun/src/adapters/tool/memory-tool-provider.ts @@ -18,7 +18,9 @@ export function buildMemoryToolProviders(store: MemoryStore | undefined): Capabi properties: { content: { type: 'string' }, scope: { type: 'string', enum: ['user', 'workspace', 'project'] }, - tags: { type: 'array', items: { type: 'string' } } + tags: { type: 'array', items: { type: 'string' } }, + ttlDays: { type: 'number', minimum: 1, description: 'Optional lifetime in days.' }, + supersedes: { type: 'string', description: 'Optional memory id replaced by this memory.' } }, required: ['content'], additionalProperties: false @@ -36,6 +38,13 @@ export function buildMemoryToolProviders(store: MemoryStore | undefined): Capabi ...(args.scope === 'project' ? { project: context.workspace } : {}), sourceThreadId: context.threadId, sourceTurnId: context.turnId, + provenance: { kind: 'user', turnId: context.turnId, origin: 'memory_create' }, + ...(typeof args.ttlDays === 'number' && Number.isFinite(args.ttlDays) && args.ttlDays > 0 + ? { ttlMs: Math.round(args.ttlDays * 24 * 60 * 60 * 1_000) } + : {}), + ...(typeof args.supersedes === 'string' && args.supersedes.trim() + ? { supersedes: args.supersedes.trim() } + : {}), tags: Array.isArray(args.tags) ? args.tags.filter((tag): tag is string => typeof tag === 'string') : [] }) } diff --git a/kun/src/adapters/tool/task-graph-tool.test.ts b/kun/src/adapters/tool/task-graph-tool.test.ts new file mode 100644 index 000000000..311c6dd02 --- /dev/null +++ b/kun/src/adapters/tool/task-graph-tool.test.ts @@ -0,0 +1,65 @@ +import { afterEach, describe, expect, it, vi } from 'vitest' +import { mkdtemp, rm } from 'node:fs/promises' +import { join } from 'node:path' +import { tmpdir } from 'node:os' +import { createTaskGraphTool } from './task-graph-tool.js' +import type { ToolHostContext } from '../../ports/tool-host.js' + +function ctx(threadId = 't1'): ToolHostContext { + return { + threadId, turnId: 'tn', workspace: '/ws', approvalPolicy: 'auto', + abortSignal: new AbortController().signal, awaitApproval: vi.fn(async () => 'allow' as const) + } +} + +describe('task_graph tool', () => { + const dirs: string[] = [] + afterEach(async () => Promise.all(dirs.splice(0).map((dir) => rm(dir, { recursive: true, force: true })))) + it('adds tasks with dependencies and reports runnable set', async () => { + const tool = createTaskGraphTool() + await tool.execute({ action: 'add', id: 'a', title: 'A' }, ctx()) + const r = await tool.execute({ action: 'add', id: 'b', title: 'B', dependsOn: ['a'] }, ctx()) + expect(r.output).toMatchObject({ runnable: ['a'] }) + }) + + it('advances the plan through start/complete and unblocks dependents', async () => { + const tool = createTaskGraphTool() + await tool.execute({ action: 'add', id: 'a', title: 'A' }, ctx()) + await tool.execute({ action: 'add', id: 'b', title: 'B', dependsOn: ['a'] }, ctx()) + await tool.execute({ action: 'start', id: 'a' }, ctx()) + const done = await tool.execute({ action: 'complete', id: 'a' }, ctx()) + expect(done.output).toMatchObject({ runnable: ['b'] }) + }) + + it('rejects a dependency cycle', async () => { + const tool = createTaskGraphTool() + await tool.execute({ action: 'add', id: 'a', title: 'A', dependsOn: ['b'] }, ctx()) + const r = await tool.execute({ action: 'add', id: 'b', title: 'B', dependsOn: ['a'] }, ctx()) + expect(r.isError).toBe(true) + expect(String((r.output as Record).error)).toContain('cycle') + }) + + it('keeps per-thread state isolated', async () => { + const tool = createTaskGraphTool() + await tool.execute({ action: 'add', id: 'a', title: 'A' }, ctx('t1')) + const other = await tool.execute({ action: 'list' }, ctx('t2')) + expect(other.output).toMatchObject({ tasks: [] }) + }) + + it('reports retry on failure with attempts remaining', async () => { + const tool = createTaskGraphTool() + await tool.execute({ action: 'add', id: 'a', title: 'A', maxAttempts: 2 }, ctx()) + await tool.execute({ action: 'start', id: 'a' }, ctx()) + const failed = await tool.execute({ action: 'fail', id: 'a', error: 'boom' }, ctx()) + expect(failed.output).toMatchObject({ retried: true }) + }) + + it('restores a thread graph from disk', async () => { + const rootDir = await mkdtemp(join(tmpdir(), 'kun-task-graph-')) + dirs.push(rootDir) + const tool = createTaskGraphTool({ rootDir }) + await tool.execute({ action: 'add', id: 'a', title: 'A' }, ctx('t1')) + const restored = await createTaskGraphTool({ rootDir }).execute({ action: 'list' }, ctx('t1')) + expect(restored.output).toMatchObject({ tasks: [{ id: 'a', title: 'A' }] }) + }) +}) diff --git a/kun/src/adapters/tool/task-graph-tool.ts b/kun/src/adapters/tool/task-graph-tool.ts new file mode 100644 index 000000000..f129fac01 --- /dev/null +++ b/kun/src/adapters/tool/task-graph-tool.ts @@ -0,0 +1,135 @@ +/** + * Agent-callable task graph (P1 #2). + * + * Lets the model plan and drive a dependency-aware task DAG instead of a flat + * todo list: add tasks with dependencies/priority/retries, ask which are + * runnable now (respecting deps + concurrency), and record start/success/ + * failure. State is kept per thread in-memory so a multi-step plan persists + * across turns within the conversation. This is an agent tool, not a core-loop + * scheduler — the model decides when to advance the plan. + */ + +import { LocalToolHost, type LocalTool } from './local-tool-host.js' +import { TaskGraph } from '../../tasks/task-graph.js' +import { createHash, randomUUID } from 'node:crypto' +import { mkdir, readFile, rename, writeFile } from 'node:fs/promises' +import { join } from 'node:path' + +export function createTaskGraphTool(options: { rootDir?: string } = {}): LocalTool { + const graphs = new Map() + const graphFor = async (threadId: string): Promise => { + let graph = graphs.get(threadId) + if (graph) return graph + graph = options.rootDir ? await loadGraph(options.rootDir, threadId) : new TaskGraph({ concurrency: 1 }) + graphs.set(threadId, graph) + return graph + } + const save = async (threadId: string, graph: TaskGraph): Promise => { + if (options.rootDir) await saveGraph(options.rootDir, threadId, graph) + } + const snapshot = (graph: TaskGraph) => ({ + tasks: graph.list().map((t) => ({ + id: t.id, title: t.title, state: t.state, dependsOn: t.dependsOn, + priority: t.priority, attempts: t.attempts, maxAttempts: t.maxAttempts, + ...(t.lastError ? { lastError: t.lastError } : {}) + })), + runnable: graph.nextRunnable().map((t) => t.id), + complete: graph.isComplete() + }) + + return LocalToolHost.defineTool({ + name: 'task_graph', + description: + 'Plan and drive a dependency-aware task graph for this thread. actions: ' + + '"add" (id,title,dependsOn?,priority?,maxAttempts?), "list", "next" (runnable now), ' + + '"start" (id), "complete" (id), "fail" (id,error), "pause"/"resume"/"cancel" (id), ' + + '"set_concurrency" (concurrency). Tasks become runnable only when their dependencies succeed.', + inputSchema: { + type: 'object', + properties: { + action: { type: 'string', enum: ['add', 'list', 'next', 'start', 'complete', 'fail', 'pause', 'resume', 'cancel', 'set_concurrency'] }, + id: { type: 'string' }, + title: { type: 'string' }, + dependsOn: { type: 'array', items: { type: 'string' } }, + priority: { type: 'number' }, + maxAttempts: { type: 'number' }, + error: { type: 'string' }, + concurrency: { type: 'number' } + }, + required: ['action'], + additionalProperties: false + }, + policy: 'auto', + execute: async (args, context) => { + const graph = await graphFor(context.threadId) + const action = typeof args.action === 'string' ? args.action : '' + const id = typeof args.id === 'string' ? args.id.trim() : '' + try { + switch (action) { + case 'add': { + if (!id || typeof args.title !== 'string' || !args.title.trim()) { + return { output: { error: 'id and title are required for add' }, isError: true } + } + graph.add({ + id, + title: args.title.trim(), + ...(Array.isArray(args.dependsOn) ? { dependsOn: args.dependsOn.filter((d): d is string => typeof d === 'string') } : {}), + ...(typeof args.priority === 'number' ? { priority: args.priority } : {}), + ...(typeof args.maxAttempts === 'number' ? { maxAttempts: args.maxAttempts } : {}) + }) + await save(context.threadId, graph) + return { output: { action, added: id, ...snapshot(graph) } } + } + case 'list': + case 'next': + return { output: { action, ...snapshot(graph) } } + case 'start': + graph.markRunning(id) + await save(context.threadId, graph) + return { output: { action, id, ...snapshot(graph) } } + case 'complete': + graph.markSucceeded(id) + await save(context.threadId, graph) + return { output: { action, id, ...snapshot(graph) } } + case 'fail': { + const outcome = graph.markFailed(id, typeof args.error === 'string' ? args.error : 'unspecified failure') + await save(context.threadId, graph) + return { output: { action, id, retried: outcome.retried, ...snapshot(graph) } } + } + case 'pause': graph.pause(id); await save(context.threadId, graph); return { output: { action, id, ...snapshot(graph) } } + case 'resume': graph.resume(id); await save(context.threadId, graph); return { output: { action, id, ...snapshot(graph) } } + case 'cancel': graph.cancel(id); await save(context.threadId, graph); return { output: { action, id, ...snapshot(graph) } } + case 'set_concurrency': + graph.setConcurrency(typeof args.concurrency === 'number' ? args.concurrency : 1) + await save(context.threadId, graph) + return { output: { action, ...snapshot(graph) } } + default: + return { output: { error: `unknown action: ${action}` }, isError: true } + } + } catch (error) { + return { output: { action, error: error instanceof Error ? error.message : String(error) }, isError: true } + } + } + }) +} + +async function loadGraph(rootDir: string, threadId: string): Promise { + try { + return TaskGraph.fromJSON(JSON.parse(await readFile(graphPath(rootDir, threadId), 'utf8'))) + } catch (error) { + if ((error as NodeJS.ErrnoException).code === 'ENOENT') return new TaskGraph({ concurrency: 1 }) + throw error + } +} + +async function saveGraph(rootDir: string, threadId: string, graph: TaskGraph): Promise { + await mkdir(rootDir, { recursive: true }) + const target = graphPath(rootDir, threadId) + const temporary = `${target}.${process.pid}.${randomUUID()}.tmp` + await writeFile(temporary, JSON.stringify(graph.toJSON(), null, 2), 'utf8') + await rename(temporary, target) +} + +function graphPath(rootDir: string, threadId: string): string { + return join(rootDir, `${createHash('sha256').update(threadId).digest('hex')}.json`) +} diff --git a/kun/src/adapters/workspace/local-workspace-inspector.ts b/kun/src/adapters/workspace/local-workspace-inspector.ts index 5ddab0fe9..9bdf866de 100644 --- a/kun/src/adapters/workspace/local-workspace-inspector.ts +++ b/kun/src/adapters/workspace/local-workspace-inspector.ts @@ -1,6 +1,6 @@ import { existsSync } from 'node:fs' import { resolve } from 'node:path' -import { execFile } from 'node:child_process' +import { execFile, type ExecFileOptions } from 'node:child_process' import { promisify } from 'node:util' import type { WorkspaceInspector } from '../../ports/workspace-inspector.js' import type { WorkspaceStatus } from '../../contracts/workspace.js' @@ -16,12 +16,17 @@ const execFileAsync = promisify(execFile) export class LocalWorkspaceInspector implements WorkspaceInspector { private readonly exec: ( file: string, - args: string[] + args: string[], + options?: ExecFileOptions ) => Promise<{ stdout: string; stderr: string }> constructor( options: { - exec?: (file: string, args: string[]) => Promise<{ stdout: string; stderr: string }> + exec?: ( + file: string, + args: string[], + options?: ExecFileOptions + ) => Promise<{ stdout: string; stderr: string }> } = {} ) { this.exec = options.exec ?? execFileAsync @@ -80,7 +85,7 @@ export class LocalWorkspaceInspector implements WorkspaceInspector { } private async runGit(cwd: string, args: string[]): Promise { - const { stdout } = await this.exec('git', args) + const { stdout } = await this.exec('git', args, { cwd }) return stdout.trim() } } diff --git a/kun/src/artifacts/artifact-store.test.ts b/kun/src/artifacts/artifact-store.test.ts new file mode 100644 index 000000000..23dd12460 --- /dev/null +++ b/kun/src/artifacts/artifact-store.test.ts @@ -0,0 +1,192 @@ +import { describe, expect, it } from 'vitest' +import { mkdtemp, rm } from 'node:fs/promises' +import { tmpdir } from 'node:os' +import { join } from 'node:path' +import { FileArtifactStore, InMemoryArtifactStore, type ArtifactStore } from './artifact-store.js' + +function runStoreContract(name: string, make: () => Promise<{ store: ArtifactStore; cleanup?: () => Promise }>) { + describe(name, () => { + it('stores content and returns a bounded summary + stable id', async () => { + const { store, cleanup } = await make() + try { + const big = 'x'.repeat(10_000) + const result = await store.put({ content: big, maxInlineChars: 200, source: 'mcp', origin: 'docs/lookup' }) + expect(result.meta.byteSize).toBe(10_000) + expect(result.summary.truncated).toBe(true) + expect(result.summary.inline.length).toBeLessThan(400) + expect(result.meta.id).toBe(result.summary.artifactId) + expect(await store.get(result.meta.id)).toBe(big) + } finally { + await cleanup?.() + } + }) + + it('dedupes identical content by hash', async () => { + const { store, cleanup } = await make() + try { + const a = await store.put({ content: 'same' }) + const b = await store.put({ content: 'same' }) + expect(a.meta.id).toBe(b.meta.id) + expect(b.deduped).toBe(true) + } finally { + await cleanup?.() + } + }) + + it('reads a line range on demand', async () => { + const { store, cleanup } = await make() + try { + const content = ['l1', 'l2', 'l3', 'l4', 'l5'].join('\n') + const { meta } = await store.put({ content }) + expect(await store.readRange(meta.id, { startLine: 2, endLine: 4 })).toBe('l2\nl3\nl4') + } finally { + await cleanup?.() + } + }) + + it('reads a byte range on demand', async () => { + const { store, cleanup } = await make() + try { + const { meta } = await store.put({ content: 'abcdefghij' }) + expect(await store.readRange(meta.id, { offset: 2, length: 3 })).toBe('cde') + } finally { + await cleanup?.() + } + }) + + it('returns null for an unknown id', async () => { + const { store, cleanup } = await make() + try { + expect(await store.get('art_missing')).toBeNull() + expect(await store.readRange('art_missing', {})).toBeNull() + expect(await store.stat('art_missing')).toBeNull() + } finally { + await cleanup?.() + } + }) + + it('records source metadata', async () => { + const { store, cleanup } = await make() + try { + const { meta } = await store.put({ content: 'hi', source: 'web', origin: 'web_fetch' }) + const stat = await store.stat(meta.id) + expect(stat).toMatchObject({ source: 'web', origin: 'web_fetch' }) + } finally { + await cleanup?.() + } + }) + }) +} + +runStoreContract('InMemoryArtifactStore', async () => ({ + store: new InMemoryArtifactStore(() => '2026-06-29T00:00:00.000Z') +})) + +runStoreContract('FileArtifactStore', async () => { + const dir = await mkdtemp(join(tmpdir(), 'kun-artifacts-')) + return { + store: new FileArtifactStore(dir, () => '2026-06-29T00:00:00.000Z'), + cleanup: () => rm(dir, { recursive: true, force: true }) + } +}) + +describe('FileArtifactStore streaming reads', () => { + it('seeks a byte range and a line window from a large artifact without loading it whole', async () => { + const dir = await mkdtemp(join(tmpdir(), 'kun-artifacts-')) + try { + const store = new FileArtifactStore(dir, () => 't0') + const lines = Array.from({ length: 100_000 }, (_, i) => `line ${i + 1}`) + const content = lines.join('\n') + const { meta } = await store.put({ content }) + // Line window stops early — only the selected lines come back. + expect(await store.readRange(meta.id, { startLine: 5, endLine: 7 })).toBe('line 5\nline 6\nline 7') + // Byte range seek. + expect(await store.readRange(meta.id, { offset: 0, length: 6 })).toBe('line 1') + } finally { + await rm(dir, { recursive: true, force: true }) + } + }) + + it('returns the trailing line when no final newline and the range covers it', async () => { + const dir = await mkdtemp(join(tmpdir(), 'kun-artifacts-')) + try { + const store = new FileArtifactStore(dir, () => 't0') + const { meta } = await store.put({ content: 'a\nb\nc' }) + expect(await store.readRange(meta.id, { startLine: 3, endLine: 3 })).toBe('c') + } finally { + await rm(dir, { recursive: true, force: true }) + } + }) + + it('stitches a multibyte UTF-8 char split across a 64KiB read boundary (P2-04)', async () => { + const dir = await mkdtemp(join(tmpdir(), 'kun-artifacts-')) + try { + const store = new FileArtifactStore(dir, () => 't0') + // Pad so a 3-byte char (世) straddles the 65536-byte chunk boundary, then + // read the line that contains it; without incremental decoding the char + // would surface as replacement characters. + const head = 'a'.repeat(65_535) + const content = `${head}世界\nsecond line` + const { meta } = await store.put({ content }) + const firstLine = await store.readRange(meta.id, { startLine: 1, endLine: 1 }) + expect(firstLine).toBe(`${head}世界`) + expect(firstLine).not.toContain('\uFFFD') + } finally { + await rm(dir, { recursive: true, force: true }) + } + }) +}) + +describe('readArtifactBounded', () => { + it('pages UTF-8 text without replacement characters or skipped bytes', async () => { + const { readArtifactBounded } = await import('./artifact-store.js') + const store = new InMemoryArtifactStore(() => 't0') + const { meta } = await store.put({ content: '中文测试' }) + const pages: string[] = [] + let offset = 0 + while (offset < meta.byteSize) { + const page = await readArtifactBounded(store, meta.id, meta, { offset, length: 2 }) + expect(page).not.toBeNull() + pages.push(page!.content) + expect(page!.content).not.toContain('�') + if (!page!.nextOffset) break + expect(page!.nextOffset).toBeGreaterThan(offset) + offset = page!.nextOffset + } + expect(pages.join('')).toBe('中文测试') + }) + + it('clamps a no-range read to the byte cap and returns a cursor', async () => { + const { readArtifactBounded, ARTIFACT_MAX_READ_BYTES } = await import('./artifact-store.js') + const store = new InMemoryArtifactStore(() => 't0') + const big = 'x'.repeat(ARTIFACT_MAX_READ_BYTES + 5_000) + const { meta } = await store.put({ content: big }) + const result = await readArtifactBounded(store, meta.id, meta, {}) + expect(result).not.toBeNull() + expect(Buffer.byteLength(result!.content, 'utf8')).toBe(ARTIFACT_MAX_READ_BYTES) + expect(result!.truncated).toBe(true) + expect(result!.nextOffset).toBe(ARTIFACT_MAX_READ_BYTES) + }) + + it('clamps an oversized line range to the line cap and returns a line cursor', async () => { + const { readArtifactBounded, ARTIFACT_MAX_READ_LINES } = await import('./artifact-store.js') + const store = new InMemoryArtifactStore(() => 't0') + const content = Array.from({ length: ARTIFACT_MAX_READ_LINES + 500 }, (_, i) => `l${i + 1}`).join('\n') + const { meta } = await store.put({ content }) + const result = await readArtifactBounded(store, meta.id, meta, { startLine: 1, endLine: ARTIFACT_MAX_READ_LINES + 500 }) + expect(result).not.toBeNull() + expect(result!.range.endLine).toBe(ARTIFACT_MAX_READ_LINES) + expect(result!.truncated).toBe(true) + expect(result!.nextStartLine).toBe(ARTIFACT_MAX_READ_LINES + 1) + }) + + it('reports not-truncated and no cursor when the whole artifact fits', async () => { + const { readArtifactBounded } = await import('./artifact-store.js') + const store = new InMemoryArtifactStore(() => 't0') + const { meta } = await store.put({ content: 'small content' }) + const result = await readArtifactBounded(store, meta.id, meta, {}) + expect(result!.content).toBe('small content') + expect(result!.truncated).toBe(false) + expect(result!.nextOffset).toBeUndefined() + }) +}) diff --git a/kun/src/artifacts/artifact-store.ts b/kun/src/artifacts/artifact-store.ts new file mode 100644 index 000000000..df6fd760f --- /dev/null +++ b/kun/src/artifacts/artifact-store.ts @@ -0,0 +1,384 @@ +/** + * Content-addressed Artifact Store (P0 #5). + * + * A unified mechanism for large tool results — MCP payloads, browser output, + * big JSON, remote logs, attachments — so the model only ever sees a bounded + * summary + a stable artifact id, and can fetch specific byte/line ranges on + * demand. Content is addressed by hash, so identical results dedupe. Two + * implementations: an in-memory store (tests / ephemeral runtime) and a + * file-backed store (persistent, survives restart for replay/audit). + */ + +import { mkdir, readFile, readdir, rename, rm, writeFile, stat as fsStat, open as fsOpen } from 'node:fs/promises' +import { join } from 'node:path' +import { StringDecoder } from 'node:string_decoder' +import { artifactId, summarizeForModel, type ArtifactSummary } from './artifact-summary.js' + +export type ArtifactSourceKind = 'mcp' | 'web' | 'bash' | 'attachment' | 'remote-log' | 'tool' | 'other' + +export type StoredArtifactMeta = { + id: string + byteSize: number + lineCount: number + mimeType?: string + source?: ArtifactSourceKind + /** Tool / origin label, e.g. an MCP tool name or `web_fetch`. */ + origin?: string + createdAt: string +} + +export type PutArtifactInput = { + content: string + mimeType?: string + source?: ArtifactSourceKind + origin?: string + /** Inline preview budget handed to the model. Default 4000. */ + maxInlineChars?: number +} + +export type PutArtifactResult = { + meta: StoredArtifactMeta + summary: ArtifactSummary + /** True when this content already existed (deduped by hash). */ + deduped: boolean +} + +export type ReadRangeOptions = { + /** Byte offset (UTF-8) for a raw slice. */ + offset?: number + length?: number + /** 1-indexed inclusive line range (takes precedence over byte offset). */ + startLine?: number + endLine?: number +} + +export interface ArtifactStore { + put(input: PutArtifactInput): Promise + get(id: string): Promise + readRange(id: string, options: ReadRangeOptions): Promise + stat(id: string): Promise +} + +const DEFAULT_MAX_INLINE = 4_000 + +/** Hard cap on any single artifact read: bytes and lines (P0 #5). A read that + * would exceed these is clamped and a cursor is returned so the caller pages. */ +export const ARTIFACT_MAX_READ_BYTES = 1_048_576 +export const ARTIFACT_MAX_READ_LINES = 2_000 + +/** Artifact ids are content hashes; reject anything else so an id can never + * escape the store directory (path traversal) when used in a file path. */ +const ARTIFACT_ID_PATTERN = /^art_[0-9a-f]{1,64}$/ + +export function isValidArtifactId(id: string): boolean { + return ARTIFACT_ID_PATTERN.test(id) +} + +export type BoundedArtifactRead = { + content: string + range: ReadRangeOptions + /** True when more content remains beyond what was returned. */ + truncated: boolean + /** Byte cursor for the next page (byte-range reads). */ + nextOffset?: number + /** Line cursor for the next page (line-range reads). */ + nextStartLine?: number +} + +/** + * Read an artifact with a HARD upper bound (P0 #5). Any request — including one + * with no range, or a range larger than the cap — is clamped to at most + * {@link ARTIFACT_MAX_READ_BYTES} / {@link ARTIFACT_MAX_READ_LINES}, and a + * cursor (`nextOffset` / `nextStartLine`) is returned when content remains so a + * caller can page instead of pulling a multi-GB artifact into memory/context. + * Byte accounting is derived from the requested window + the artifact size (not + * the decoded string) so the cursor is exact even when a byte slice lands on a + * multibyte UTF-8 boundary. + */ +export async function readArtifactBounded( + store: ArtifactStore, + id: string, + meta: StoredArtifactMeta, + requested: ReadRangeOptions +): Promise { + const lineMode = requested.startLine !== undefined || requested.endLine !== undefined + if (lineMode) { + const startLine = Math.max(1, requested.startLine ?? 1) + const requestedEnd = requested.endLine ?? startLine + ARTIFACT_MAX_READ_LINES - 1 + const endLine = Math.min(requestedEnd, startLine + ARTIFACT_MAX_READ_LINES - 1) + const range: ReadRangeOptions = { startLine, endLine } + const content = await store.readRange(id, range) + if (content === null) return null + const truncated = endLine < meta.lineCount + return { content, range, truncated, ...(truncated ? { nextStartLine: endLine + 1 } : {}) } + } + const offset = Math.max(0, requested.offset ?? 0) + const requestedLen = requested.length ?? ARTIFACT_MAX_READ_BYTES + const length = Math.min(Math.max(0, requestedLen), ARTIFACT_MAX_READ_BYTES) + const content = await store.readRange(id, { offset, length }) + if (content === null) return null + const bytesConsumed = Buffer.byteLength(content, 'utf8') + const range: ReadRangeOptions = { offset, length: bytesConsumed } + const truncated = offset + bytesConsumed < meta.byteSize + return { content, range, truncated, ...(truncated ? { nextOffset: offset + bytesConsumed } : {}) } +} + +function buildMeta(input: PutArtifactInput, id: string, nowIso: () => string): StoredArtifactMeta { + const byteSize = Buffer.byteLength(input.content, 'utf8') + const lineCount = input.content.length === 0 ? 0 : input.content.split('\n').length + return { + id, + byteSize, + lineCount, + ...(input.mimeType ? { mimeType: input.mimeType } : {}), + ...(input.source ? { source: input.source } : {}), + ...(input.origin ? { origin: input.origin } : {}), + createdAt: nowIso() + } +} + +function sliceContent(content: string, options: ReadRangeOptions): string { + if (options.startLine !== undefined || options.endLine !== undefined) { + const lines = content.split('\n') + const start = Math.max(1, options.startLine ?? 1) + const end = Math.min(lines.length, options.endLine ?? lines.length) + if (start > end) return '' + return lines.slice(start - 1, end).join('\n') + } + if (options.offset !== undefined || options.length !== undefined) { + const buffer = Buffer.from(content, 'utf8') + const offset = Math.max(0, options.offset ?? 0) + const length = options.length !== undefined ? Math.max(0, options.length) : buffer.length - offset + return decodeUtf8Window(buffer.subarray(offset, offset + length + 3), length) + } + return content +} + +export class InMemoryArtifactStore implements ArtifactStore { + private readonly contents = new Map() + private readonly metas = new Map() + + constructor(private readonly nowIso: () => string = () => new Date().toISOString()) {} + + async put(input: PutArtifactInput): Promise { + const id = artifactId(input.content) + const summary = summarizeForModel({ + content: input.content, + maxInlineChars: input.maxInlineChars ?? DEFAULT_MAX_INLINE + }) + const deduped = this.contents.has(id) + if (!deduped) { + this.contents.set(id, input.content) + this.metas.set(id, buildMeta(input, id, this.nowIso)) + } + return { meta: this.metas.get(id)!, summary, deduped } + } + + async get(id: string): Promise { + return this.contents.get(id) ?? null + } + + async readRange(id: string, options: ReadRangeOptions): Promise { + const content = this.contents.get(id) + if (content === undefined) return null + return sliceContent(content, options) + } + + async stat(id: string): Promise { + return this.metas.get(id) ?? null + } +} + +export class FileArtifactStore implements ArtifactStore { + private ready?: Promise + + constructor( + private readonly dir: string, + private readonly nowIso: () => string = () => new Date().toISOString(), + private readonly limits: { maxTotalBytes?: number; maxArtifacts?: number } = {} + ) {} + + private async ensureDir(): Promise { + if (!this.ready) this.ready = mkdir(this.dir, { recursive: true }).then(() => undefined) + return this.ready + } + + private contentPath(id: string): string { + return join(this.dir, `${id}.bin`) + } + + private metaPath(id: string): string { + return join(this.dir, `${id}.json`) + } + + async put(input: PutArtifactInput): Promise { + await this.ensureDir() + const id = artifactId(input.content) + const summary = summarizeForModel({ + content: input.content, + maxInlineChars: input.maxInlineChars ?? DEFAULT_MAX_INLINE + }) + let deduped = true + try { + await Promise.all([fsStat(this.contentPath(id)), fsStat(this.metaPath(id))]) + } catch { + deduped = false + } + let meta: StoredArtifactMeta + if (!deduped) { + meta = buildMeta(input, id, this.nowIso) + await this.enforceQuota(meta.byteSize) + const suffix = `${process.pid}.${Date.now()}.tmp` + const contentTemporaryPath = `${this.contentPath(id)}.${suffix}` + const metaTemporaryPath = `${this.metaPath(id)}.${suffix}` + try { + await writeFile(contentTemporaryPath, input.content, 'utf8') + await writeFile(metaTemporaryPath, JSON.stringify(meta), 'utf8') + await rename(contentTemporaryPath, this.contentPath(id)) + await rename(metaTemporaryPath, this.metaPath(id)) + } finally { + await Promise.all([ + rm(contentTemporaryPath, { force: true }), + rm(metaTemporaryPath, { force: true }) + ]).catch(() => undefined) + } + } else { + meta = (await this.stat(id)) ?? buildMeta(input, id, this.nowIso) + } + return { meta, summary, deduped } + } + + async get(id: string): Promise { + if (!isValidArtifactId(id)) return null + try { + return await readFile(this.contentPath(id), 'utf8') + } catch { + return null + } + } + + async readRange(id: string, options: ReadRangeOptions): Promise { + if (!isValidArtifactId(id)) return null + if (options.startLine !== undefined || options.endLine !== undefined) { + return this.readLineRange(id, options.startLine, options.endLine) + } + if (options.offset !== undefined || options.length !== undefined) { + return this.readByteRange(id, options.offset ?? 0, options.length) + } + return this.get(id) + } + + /** True seek read of a byte window — never loads the whole artifact. */ + private async readByteRange(id: string, offset: number, length?: number): Promise { + let handle: import('node:fs/promises').FileHandle | undefined + try { + handle = await fsOpen(this.contentPath(id), 'r') + const { size } = await handle.stat() + const start = Math.max(0, Math.min(offset, size)) + const want = length !== undefined ? Math.max(0, length) : size - start + const count = Math.min(want + 3, size - start) + if (count <= 0) return '' + const buffer = Buffer.alloc(count) + await handle.read(buffer, 0, count, start) + return decodeUtf8Window(buffer, want) + } catch { + return null + } finally { + await handle?.close().catch(() => undefined) + } + } + + /** + * Stream a 1-indexed inclusive line window, stopping once `endLine` is read — + * so a small slice of a multi-GB artifact stays cheap. Only the selected + * lines are buffered. + */ + private async readLineRange(id: string, startLine?: number, endLine?: number): Promise { + const start = Math.max(1, startLine ?? 1) + const end = endLine ?? Number.MAX_SAFE_INTEGER + if (start > end) return '' + let handle: import('node:fs/promises').FileHandle | undefined + try { + handle = await fsOpen(this.contentPath(id), 'r') + } catch { + return null + } + try { + const collected: string[] = [] + let lineNo = 1 + let pending = '' + // Decode incrementally so a multibyte UTF-8 sequence split across a 64KiB + // read boundary is buffered and stitched, not turned into replacement + // characters (P2-04). + const decoder = new StringDecoder('utf8') + const chunk = Buffer.alloc(64 * 1_024) + for (;;) { + const { bytesRead } = await handle.read(chunk, 0, chunk.length, null) + if (bytesRead === 0) break + pending += decoder.write(chunk.subarray(0, bytesRead)) + let nl = pending.indexOf('\n') + while (nl !== -1) { + const line = pending.slice(0, nl) + if (lineNo >= start && lineNo <= end) collected.push(line) + lineNo += 1 + if (lineNo > end) return collected.join('\n') + pending = pending.slice(nl + 1) + nl = pending.indexOf('\n') + } + } + // Flush any bytes the decoder was holding for an incomplete sequence. + pending += decoder.end() + // Trailing line without a final newline. + if (lineNo >= start && lineNo <= end && pending.length > 0) collected.push(pending) + return collected.join('\n') + } catch { + return null + } finally { + await handle.close().catch(() => undefined) + } + } + + async stat(id: string): Promise { + if (!isValidArtifactId(id)) return null + try { + return JSON.parse(await readFile(this.metaPath(id), 'utf8')) as StoredArtifactMeta + } catch { + return null + } + } + + private async enforceQuota(incomingBytes: number): Promise { + const maxTotalBytes = this.limits.maxTotalBytes ?? 512 * 1024 * 1024 + const maxArtifacts = this.limits.maxArtifacts ?? 2_000 + if (incomingBytes > maxTotalBytes) throw new Error(`artifact exceeds ${maxTotalBytes} byte store quota`) + const entries = await readdir(this.dir).catch(() => []) + const metas = (await Promise.all(entries + .filter((entry) => entry.endsWith('.json')) + .map(async (entry) => { + try { + return JSON.parse(await readFile(join(this.dir, entry), 'utf8')) as StoredArtifactMeta + } catch { + return null + } + }))) + .filter((meta): meta is StoredArtifactMeta => Boolean(meta && isValidArtifactId(meta.id))) + const totalBytes = metas.reduce((sum, meta) => sum + meta.byteSize, 0) + if (totalBytes + incomingBytes > maxTotalBytes || metas.length + 1 > maxArtifacts) { + throw new Error('artifact store quota exceeded; remove unneeded artifacts before retrying') + } + } +} + +function decodeUtf8Window(buffer: Buffer, requestedBytes: number): string { + const decode = (bytes: Buffer): string => { + const decoder = new StringDecoder('utf8') + return decoder.write(bytes) + } + const bounded = buffer.subarray(0, Math.min(requestedBytes, buffer.length)) + const text = decode(bounded) + if (text || bounded.length === 0 || buffer.length <= bounded.length) return text + // Very small ranges can end before the first multibyte code point completes. + // Include that one complete character so the cursor advances without losing + // bytes or looping forever; the overrun is at most three bytes. + return decode(buffer.subarray(0, Math.min(requestedBytes + 3, buffer.length))) +} diff --git a/kun/src/artifacts/artifact-summary.test.ts b/kun/src/artifacts/artifact-summary.test.ts new file mode 100644 index 000000000..376ffe237 --- /dev/null +++ b/kun/src/artifacts/artifact-summary.test.ts @@ -0,0 +1,42 @@ +import { describe, expect, it } from 'vitest' +import { artifactId, summarizeForModel } from './artifact-summary.js' + +describe('artifactId', () => { + it('is content-addressed and stable', () => { + expect(artifactId('hello')).toBe(artifactId('hello')) + expect(artifactId('hello')).not.toBe(artifactId('world')) + expect(artifactId('hello')).toMatch(/^art_[0-9a-f]{24}$/) + }) +}) + +describe('summarizeForModel', () => { + it('returns content verbatim when within budget', () => { + const result = summarizeForModel({ content: 'short output', maxInlineChars: 100 }) + expect(result.truncated).toBe(false) + expect(result.inline).toBe('short output') + expect(result.lineCount).toBe(1) + }) + + it('truncates with head + tail + reference marker when oversized', () => { + const content = `${'H'.repeat(500)}\n${'T'.repeat(500)}` + const result = summarizeForModel({ content, maxInlineChars: 200 }) + expect(result.truncated).toBe(true) + expect(result.inline).toContain(`[artifact ${result.artifactId}:`) + expect(result.inline).toContain('elided') + // Keeps a head and a tail excerpt. + expect(result.inline.startsWith('H')).toBe(true) + expect(result.inline.trimEnd().endsWith('T')).toBe(true) + // The byte size reflects the full content, not the preview. + expect(result.byteSize).toBe(Buffer.byteLength(content, 'utf8')) + }) + + it('keeps the preview within roughly the requested budget', () => { + const result = summarizeForModel({ content: 'x'.repeat(10_000), maxInlineChars: 300 }) + // head + marker + tail; allow the marker line + newlines as overhead. + expect(result.inline.length).toBeLessThanOrEqual(400) + }) + + it('counts lines on the full content', () => { + expect(summarizeForModel({ content: 'a\nb\nc', maxInlineChars: 100 }).lineCount).toBe(3) + }) +}) diff --git a/kun/src/artifacts/artifact-summary.ts b/kun/src/artifacts/artifact-summary.ts new file mode 100644 index 000000000..3cb9f9c53 --- /dev/null +++ b/kun/src/artifacts/artifact-summary.ts @@ -0,0 +1,77 @@ +/** + * Content-addressed artifact summarization (P0 #5). + * + * Bash output already has truncation, but MCP results, browser output, large + * JSON, and attachments do not share a mechanism. Large results should be + * stored by content hash and the model should see only a bounded preview plus + * a stable reference, fetching specific slices on demand. This module is the + * pure core: a content-address id + a head/tail-bounded preview with an + * explicit reference marker. The byte store itself is a thin layer on top. + */ + +import { createHash } from 'node:crypto' + +export function artifactId(content: string | Buffer): string { + const hash = createHash('sha256').update(content).digest('hex') + return `art_${hash.slice(0, 24)}` +} + +export type ArtifactSummary = { + artifactId: string + byteSize: number + lineCount: number + /** The bounded preview handed to the model. */ + inline: string + /** True when `inline` is a head/tail excerpt rather than the full content. */ + truncated: boolean +} + +/** + * Produce a model-facing summary of a (possibly huge) text result. When the + * content fits within `maxInlineChars` it is returned verbatim. Otherwise a + * head + tail excerpt is returned with a reference marker telling the model how + * many bytes/lines were elided and how to fetch the full artifact by id. + */ +export function summarizeForModel(input: { + content: string + maxInlineChars: number + /** Fraction of the budget devoted to the head excerpt (rest is tail). Default 0.7. */ + headRatio?: number +}): ArtifactSummary { + const content = input.content + const byteSize = Buffer.byteLength(content, 'utf8') + const lineCount = content.length === 0 ? 0 : content.split('\n').length + const id = artifactId(content) + const maxInline = Math.max(0, input.maxInlineChars) + + if (content.length <= maxInline) { + return { artifactId: id, byteSize, lineCount, inline: content, truncated: false } + } + + const headRatio = clampRatio(input.headRatio ?? 0.7) + // Reserve room for the marker so the final inline stays within budget. + const markerPlaceholder = referenceMarker(id, 0, 0) + const budget = Math.max(0, maxInline - markerPlaceholder.length) + const headChars = Math.floor(budget * headRatio) + const tailChars = Math.max(0, budget - headChars) + const head = content.slice(0, headChars) + const tail = tailChars > 0 ? content.slice(content.length - tailChars) : '' + const omittedChars = content.length - head.length - tail.length + const omittedLines = Math.max(0, lineCount - countLines(head) - countLines(tail)) + const marker = referenceMarker(id, omittedChars, omittedLines) + const inline = tail ? `${head}\n${marker}\n${tail}` : `${head}\n${marker}` + return { artifactId: id, byteSize, lineCount, inline, truncated: true } +} + +function referenceMarker(id: string, omittedChars: number, omittedLines: number): string { + return `[artifact ${id}: ${omittedChars} char(s) / ${omittedLines} line(s) elided — fetch the full artifact or a specific range by id]` +} + +function countLines(text: string): number { + return text.length === 0 ? 0 : text.split('\n').length +} + +function clampRatio(value: number): number { + if (!Number.isFinite(value)) return 0.7 + return Math.min(0.95, Math.max(0.05, value)) +} diff --git a/kun/src/attachments/attachment-store.ts b/kun/src/attachments/attachment-store.ts index 3543da8e2..f0e7bc8ce 100644 --- a/kun/src/attachments/attachment-store.ts +++ b/kun/src/attachments/attachment-store.ts @@ -5,6 +5,8 @@ import type { AttachmentsCapabilityConfig } from '../contracts/capabilities.js' import type { AttachmentDiagnostics, AttachmentMetadata, AttachmentTextFallback } from '../contracts/attachments.js' import { AttachmentMetadata as AttachmentMetadataSchema } from '../contracts/attachments.js' +const ATTACHMENT_ID_PATTERN = /^att_[0-9a-f]{24}$/ + export type AttachmentContent = AttachmentMetadata & { data: Buffer } @@ -14,6 +16,8 @@ export interface AttachmentStore { name: string data: Buffer mimeType?: string + documentText?: string + pageCount?: number localFilePath?: string textFallback?: AttachmentTextFallback threadId?: string @@ -41,6 +45,8 @@ export class FileAttachmentStore implements AttachmentStore { name: string data: Buffer mimeType?: string + documentText?: string + pageCount?: number localFilePath?: string textFallback?: AttachmentTextFallback threadId?: string @@ -48,14 +54,7 @@ export class FileAttachmentStore implements AttachmentStore { }): Promise { await mkdir(this.options.rootDir, { recursive: true }) const image = detectImage(input.data) - if (!image) throw new Error('unsupported image MIME type') - if (input.mimeType && input.mimeType !== image.mimeType) throw new Error('declared MIME type does not match image content') - if (!this.options.config.allowedMimeTypes.includes(image.mimeType)) throw new Error(`image MIME type is not allowed: ${image.mimeType}`) - if (input.data.byteLength > this.options.config.maxImageBytes) throw new Error(`image exceeds ${this.options.config.maxImageBytes} byte limit`) - const maxDimension = Math.max(image.width ?? 0, image.height ?? 0) - if (maxDimension > this.options.config.maxImageDimension) { - throw new Error(`image exceeds ${this.options.config.maxImageDimension}px dimension limit`) - } + const descriptor = image ? this.describeImage(image, input) : this.describeDocument(input) if (input.textFallback) validateTextFallback(input.textFallback, this.options.config) const hash = createHash('sha256').update(input.data).digest('hex') const id = `att_${hash.slice(0, 24)}` @@ -66,8 +65,13 @@ export class FileAttachmentStore implements AttachmentStore { if (existing) { const next = mergeScope({ ...existing, + kind: descriptor.kind, + mimeType: descriptor.mimeType, ...(input.localFilePath ? { localFilePath: input.localFilePath } : {}), ...(input.textFallback ? { textFallback: input.textFallback } : {}), + ...(descriptor.documentText !== undefined ? { documentText: descriptor.documentText } : {}), + ...(descriptor.pageCount ? { pageCount: descriptor.pageCount } : {}), + ...(descriptor.truncated !== undefined ? { truncated: descriptor.truncated } : {}), updatedAt: now }, input) await writeFile(contentPath, input.data) @@ -77,11 +81,15 @@ export class FileAttachmentStore implements AttachmentStore { const metadata: AttachmentMetadata = AttachmentMetadataSchema.parse(mergeScope({ id, name: input.name, - mimeType: image.mimeType, + kind: descriptor.kind, + mimeType: descriptor.mimeType, byteSize: input.data.byteLength, hash, - ...(image.width ? { width: image.width } : {}), - ...(image.height ? { height: image.height } : {}), + ...(descriptor.width ? { width: descriptor.width } : {}), + ...(descriptor.height ? { height: descriptor.height } : {}), + ...(descriptor.documentText !== undefined ? { documentText: descriptor.documentText } : {}), + ...(descriptor.pageCount ? { pageCount: descriptor.pageCount } : {}), + ...(descriptor.truncated !== undefined ? { truncated: descriptor.truncated } : {}), ...(input.localFilePath ? { localFilePath: input.localFilePath } : {}), ...(input.textFallback ? { textFallback: input.textFallback } : {}), threadIds: [], @@ -94,7 +102,51 @@ export class FileAttachmentStore implements AttachmentStore { return metadata } + private describeImage( + image: { mimeType: string; width?: number; height?: number }, + input: { data: Buffer; mimeType?: string } + ): AttachmentDescriptor { + if (input.mimeType && input.mimeType !== image.mimeType) throw new Error('declared MIME type does not match image content') + if (!this.options.config.allowedMimeTypes.includes(image.mimeType)) throw new Error(`image MIME type is not allowed: ${image.mimeType}`) + if (input.data.byteLength > this.options.config.maxImageBytes) throw new Error(`image exceeds ${this.options.config.maxImageBytes} byte limit`) + const maxDimension = Math.max(image.width ?? 0, image.height ?? 0) + if (maxDimension > this.options.config.maxImageDimension) { + throw new Error(`image exceeds ${this.options.config.maxImageDimension}px dimension limit`) + } + return { kind: 'image', mimeType: image.mimeType, width: image.width, height: image.height } + } + + private describeDocument(input: { + data: Buffer + mimeType?: string + documentText?: string + pageCount?: number + }): AttachmentDescriptor { + const mimeType = resolveDocumentMimeType(input) + const allowed = this.options.config.allowedDocumentMimeTypes + if (!mimeType || !allowed.includes(mimeType)) { + throw new Error(`unsupported attachment type (expected an image or an allowed document, got ${mimeType ?? input.mimeType ?? 'unknown'})`) + } + if (input.data.byteLength > this.options.config.maxDocumentBytes) { + throw new Error(`document exceeds ${this.options.config.maxDocumentBytes} byte limit`) + } + const rawText = input.documentText ?? decodeTextDocument(mimeType, input.data) + if (rawText === undefined) { + throw new Error(`document text is required for ${mimeType} attachments`) + } + const limit = this.options.config.maxDocumentTextChars + const truncated = rawText.length > limit + return { + kind: 'document', + mimeType, + documentText: truncated ? rawText.slice(0, limit) : rawText, + ...(input.pageCount ? { pageCount: input.pageCount } : {}), + ...(truncated ? { truncated: true } : {}) + } + } + async get(id: string): Promise { + if (!ATTACHMENT_ID_PATTERN.test(id)) return null try { return AttachmentMetadataSchema.parse(JSON.parse(await readFile(this.metadataPath(id), 'utf8'))) } catch { @@ -103,6 +155,7 @@ export class FileAttachmentStore implements AttachmentStore { } async resolveContent(id: string, scope: { threadId?: string; workspace?: string }): Promise { + if (!ATTACHMENT_ID_PATTERN.test(id)) throw new Error(`invalid attachment id: ${id}`) const metadata = await this.get(id) if (!metadata) throw new Error(`attachment not found: ${id}`) if (!isAuthorized(metadata, scope)) throw new Error(`attachment is not authorized for this turn: ${id}`) @@ -183,6 +236,28 @@ function validateTextFallback(fallback: AttachmentTextFallback, config: Attachme } } +type AttachmentDescriptor = { + kind: 'image' | 'document' + mimeType: string + width?: number + height?: number + documentText?: string + pageCount?: number + truncated?: boolean +} + +function resolveDocumentMimeType(input: { data: Buffer; mimeType?: string }): string | undefined { + if (input.data.length >= 5 && input.data.subarray(0, 5).toString('ascii') === '%PDF-') { + return 'application/pdf' + } + return input.mimeType?.trim().toLowerCase() || undefined +} + +function decodeTextDocument(mimeType: string, data: Buffer): string | undefined { + if (!mimeType.startsWith('text/') && mimeType !== 'application/json') return undefined + return data.toString('utf8').replace(/^\uFEFF/, '') +} + export function detectImage(buffer: Buffer): { mimeType: string; width?: number; height?: number } | null { if (buffer.length >= 24 && buffer[0] === 0x89 && buffer[1] === 0x50 && buffer[2] === 0x4e && buffer[3] === 0x47) { return { mimeType: 'image/png', width: buffer.readUInt32BE(16), height: buffer.readUInt32BE(20) } diff --git a/kun/src/benchmark/replay-benchmark.test.ts b/kun/src/benchmark/replay-benchmark.test.ts index 9bf69d261..a473368c0 100644 --- a/kun/src/benchmark/replay-benchmark.test.ts +++ b/kun/src/benchmark/replay-benchmark.test.ts @@ -2,6 +2,7 @@ import { describe, expect, it } from 'vitest' import type { RuntimeEvent } from '../contracts/events.js' import { compareReplayReports, + evaluateReplayQuality, ReplaySuiteSchema, runReplaySuite, SseMessageDecoder, @@ -181,6 +182,91 @@ describe('replay benchmark', () => { })).toThrow('duplicate replay task id') }) + it('scores required output, changed files, forbidden behavior, and cost', () => { + const task = ReplaySuiteSchema.parse({ + version: 1, + name: 'quality-suite', + tasks: [{ + id: 'quality', + prompt: 'fix the pool', + expect: { + requiredOutputs: ['poolSize'], + expectedChangedFiles: ['src/db.ts'], + forbiddenBehaviors: ['force push'], + maxCostUsd: 0.01 + } + }] + }).tasks[0]! + const events: ObservedReplayEvent[] = [ + observed({ + kind: 'item_completed', + seq: 1, + timestamp: '2026-06-29T00:00:00.000Z', + threadId: 'thread_1', + turnId: 'turn_1', + item: { + ...itemBase('tool_call'), + kind: 'tool_call', + toolName: 'edit', + callId: 'call_edit', + toolKind: 'file_change', + arguments: { path: 'src/db.ts' } + } + } as RuntimeEvent, 10), + observed({ + kind: 'item_completed', + seq: 2, + timestamp: '2026-06-29T00:00:00.010Z', + threadId: 'thread_1', + turnId: 'turn_1', + item: { ...itemBase('assistant_text'), kind: 'assistant_text', text: 'Updated poolSize safely.' } + } as RuntimeEvent, 20) + ] + + const quality = evaluateReplayQuality(task, { ...replayRun('passed', 10, 20, 0.8).metrics, costUsd: 0.005 }, events) + + expect(quality).toMatchObject({ score: 1, passed: true, violations: [] }) + expect(quality.breakdown.map((entry) => entry.dimension)).toEqual([ + 'files', + 'forbidden', + 'outputs', + 'cost' + ]) + }) + + it('hard-fails replay quality when a forbidden behavior is observed', () => { + const task = ReplaySuiteSchema.parse({ + version: 1, + name: 'unsafe-suite', + tasks: [{ + id: 'unsafe', + prompt: 'publish changes', + expect: { forbiddenBehaviors: ['force push'] } + }] + }).tasks[0]! + const events = [observed({ + kind: 'item_completed', + seq: 1, + timestamp: '2026-06-29T00:00:00.000Z', + threadId: 'thread_1', + turnId: 'turn_1', + item: { + ...itemBase('tool_call'), + kind: 'tool_call', + toolName: 'bash', + callId: 'call_bash', + toolKind: 'command_execution', + arguments: { command: 'force push origin main' } + } + } as RuntimeEvent, 10)] + + const quality = evaluateReplayQuality(task, replayRun('passed', 10, 20, 0.8).metrics, events) + + expect(quality.score).toBe(0) + expect(quality.passed).toBe(false) + expect(quality.violations.join(' ')).toContain('force push') + }) + it('fails runs that do not use any required investigation tool', async () => { const fetchImpl: typeof fetch = async (input, init = {}) => { const url = new URL(String(input)) @@ -221,7 +307,10 @@ describe('replay benchmark', () => { tasks: [{ id: 'no-tool', prompt: 'answer from memory', - expect: { requiredAnyTools: ['read', 'grep', 'find', 'ls'] } + expect: { + requiredAnyTools: ['read', 'grep', 'find', 'ls'], + requiredOutputs: ['inspection complete'] + } }] }, { baseUrl: 'http://127.0.0.1:18899', @@ -232,6 +321,8 @@ describe('replay benchmark', () => { expect(report.runs[0]?.status).toBe('failed') expect(report.runs[0]?.failureReasons).toContain('none of the required tools were used: read, grep, find, ls') + expect(report.runs[0]?.failureReasons).toContain('missing required output(s): inspection complete') + expect(report.runs[0]?.quality?.passed).toBe(false) }) it('interrupts timed-out turns before deleting replay threads', async () => { diff --git a/kun/src/benchmark/replay-benchmark.ts b/kun/src/benchmark/replay-benchmark.ts index f16bb809b..d1d50f1fc 100644 --- a/kun/src/benchmark/replay-benchmark.ts +++ b/kun/src/benchmark/replay-benchmark.ts @@ -9,8 +9,12 @@ const ReplayExpectationSchema = z.object({ minAssistantChars: z.number().int().nonnegative().default(1), requiredTools: z.array(z.string().min(1)).default([]), requiredAnyTools: z.array(z.string().min(1)).default([]), + requiredOutputs: z.array(z.string().min(1)).default([]), + forbiddenBehaviors: z.array(z.string().min(1)).default([]), + expectedChangedFiles: z.array(z.string().min(1)).default([]), maxErrorEvents: z.number().int().nonnegative().default(0), - maxTotalMs: z.number().int().positive().optional() + maxTotalMs: z.number().int().positive().optional(), + maxCostUsd: z.number().nonnegative().optional() }).strict() const ReplayTaskSchema = z.object({ @@ -91,9 +95,24 @@ export type ReplayRunResult = { status: 'passed' | 'failed' | 'timeout' | 'error' failureReasons: string[] metrics: ReplayRunMetrics + quality?: ReplayQualityResult error?: string } +export type ReplayQualityDimension = { + dimension: 'files' | 'forbidden' | 'outputs' | 'cost' + score: number + weight: number + detail: string +} + +export type ReplayQualityResult = { + score: number + passed: boolean + violations: string[] + breakdown: ReplayQualityDimension[] +} + export type ReplayReportSummary = { runCount: number passed: number @@ -284,7 +303,11 @@ async function runReplayTask(input: { collected.elapsedMs, after.memoryUsage?.peakRssBytes ) - const failureReasons = replayExpectationFailures(task, collected.timedOut, metrics, collected.events) + const quality = evaluateReplayQuality(task, metrics, collected.events) + const failureReasons = [ + ...replayExpectationFailures(task, collected.timedOut, metrics, collected.events), + ...quality.violations + ] return { id: runId, taskId: task.id, @@ -294,7 +317,8 @@ async function runReplayTask(input: { turnId, status: collected.timedOut ? 'timeout' : failureReasons.length > 0 ? 'failed' : 'passed', failureReasons, - metrics + metrics, + quality } } catch (error) { shouldInterrupt = turnId !== undefined @@ -657,6 +681,134 @@ function replayExpectationFailures( return failures } +export function evaluateReplayQuality( + task: ReplayTask, + metrics: ReplayRunMetrics, + events: ObservedReplayEvent[] +): ReplayQualityResult { + const breakdown: ReplayQualityDimension[] = [] + const violations: string[] = [] + const observation = replayQualityObservation(events) + + if (task.expect.expectedChangedFiles.length > 0) { + const expected = uniqueNormalizedPaths(task.expect.expectedChangedFiles) + const actual = uniqueNormalizedPaths(observation.changedFiles) + const score = jaccard(expected, actual) + const missing = expected.filter((path) => !actual.includes(path)) + if (missing.length > 0) violations.push(`missing expected changed file(s): ${missing.join(', ')}`) + breakdown.push({ + dimension: 'files', + score, + weight: 2, + detail: `${Math.round(score * 100)}% changed-file overlap` + }) + } + + let hardFail = false + if (task.expect.forbiddenBehaviors.length > 0) { + const haystack = `${observation.behaviors.join('\n')}\n${observation.finalOutput}`.toLowerCase() + const hits = task.expect.forbiddenBehaviors.filter((value) => haystack.includes(value.toLowerCase())) + if (hits.length > 0) { + hardFail = true + violations.push(`forbidden behavior(s) detected: ${hits.join(', ')}`) + } + breakdown.push({ + dimension: 'forbidden', + score: hits.length === 0 ? 1 : 0, + weight: 3, + detail: hits.length === 0 ? 'none detected' : hits.join(', ') + }) + } + + if (task.expect.requiredOutputs.length > 0) { + const output = observation.finalOutput.toLowerCase() + const missing = task.expect.requiredOutputs.filter((value) => !output.includes(value.toLowerCase())) + const score = 1 - missing.length / task.expect.requiredOutputs.length + if (missing.length > 0) violations.push(`missing required output(s): ${missing.join(', ')}`) + breakdown.push({ + dimension: 'outputs', + score, + weight: 2, + detail: `${task.expect.requiredOutputs.length - missing.length}/${task.expect.requiredOutputs.length} present` + }) + } + + if (task.expect.maxCostUsd !== undefined) { + const withinBudget = metrics.costUsd <= task.expect.maxCostUsd + const score = withinBudget || metrics.costUsd === 0 + ? 1 + : Math.max(0, task.expect.maxCostUsd / metrics.costUsd) + if (!withinBudget) { + violations.push(`cost $${metrics.costUsd.toFixed(4)} exceeds $${task.expect.maxCostUsd.toFixed(4)}`) + } + breakdown.push({ dimension: 'cost', score, weight: 1, detail: `$${metrics.costUsd.toFixed(4)}` }) + } + + const totalWeight = breakdown.reduce((total, item) => total + item.weight, 0) + const weightedScore = totalWeight === 0 + ? 1 + : breakdown.reduce((total, item) => total + item.score * item.weight, 0) / totalWeight + return { + score: hardFail ? 0 : weightedScore, + passed: violations.length === 0, + violations, + breakdown + } +} + +function replayQualityObservation(events: ObservedReplayEvent[]): { + finalOutput: string + behaviors: string[] + changedFiles: string[] +} { + const assistantText = new Map() + const toolCalls = new Map; toolKind: string }>() + for (const { event } of events) { + if ('item' in event && event.item.kind === 'assistant_text') { + if (event.kind === 'assistant_text_delta') { + assistantText.set(event.item.id, `${assistantText.get(event.item.id) ?? ''}${event.item.text}`) + } else { + assistantText.set(event.item.id, event.item.text) + } + } + if ('item' in event && event.item.kind === 'tool_call') { + toolCalls.set(event.item.callId, { + name: event.item.toolName, + arguments: event.item.arguments, + toolKind: event.item.toolKind + }) + } + } + const changedFiles = [...toolCalls.values()] + .filter((call) => call.toolKind === 'file_change') + .flatMap((call) => filePathsFromArguments(call.arguments)) + return { + finalOutput: [...assistantText.values()].join('\n'), + behaviors: [...toolCalls.values()].map((call) => `${call.name} ${JSON.stringify(call.arguments)}`), + changedFiles + } +} + +function filePathsFromArguments(args: Record): string[] { + return ['path', 'filePath', 'file_path', 'targetPath', 'target_path'] + .map((key) => args[key]) + .filter((value): value is string => typeof value === 'string' && value.trim().length > 0) +} + +function uniqueNormalizedPaths(paths: readonly string[]): string[] { + return [...new Set(paths.map((path) => path.trim().replace(/\\/g, '/')).filter(Boolean))] +} + +function jaccard(expected: readonly string[], actual: readonly string[]): number { + if (expected.length === 0) return 1 + const expectedSet = new Set(expected) + const actualSet = new Set(actual) + let intersection = 0 + for (const value of expectedSet) if (actualSet.has(value)) intersection += 1 + const union = new Set([...expectedSet, ...actualSet]).size + return union === 0 ? 1 : intersection / union +} + function errorReplayRun(id: string, task: ReplayTask, iteration: number, error: string): ReplayRunResult { return { id, diff --git a/kun/src/cache/cache-regression.test.ts b/kun/src/cache/cache-regression.test.ts new file mode 100644 index 000000000..61345dfc6 --- /dev/null +++ b/kun/src/cache/cache-regression.test.ts @@ -0,0 +1,81 @@ +import { describe, expect, it } from 'vitest' +import { analyzeCacheRegression, explainCacheRegression } from './cache-regression.js' + +describe('analyzeCacheRegression', () => { + it('reports no regression when the hit rate holds steady', () => { + const report = analyzeCacheRegression({ + current: 0.9, + baseline: [0.92, 0.91, 0.93], + reasons: [] + }) + expect(report.severity).toBe('none') + expect(report.explanation).toBeNull() + }) + + it('classifies a catastrophic drop as a cliff and attributes the prefix change', () => { + const report = analyzeCacheRegression({ + current: 0.05, + baseline: [0.95, 0.93, 0.94], + reasons: ['stable_prefix_changed', 'provider_cache_miss'] + }) + expect(report.severity).toBe('cliff') + expect(report.primaryReason).toBe('stable_prefix_changed') + expect(report.dropPercentagePoints).toBeGreaterThan(80) + expect(report.explanation).toContain('Cache hit rate dropped') + expect(report.explanation).toContain('stable system prefix changed') + }) + + it('classifies a mid-size drop as major and a small one as minor', () => { + expect(analyzeCacheRegression({ current: 0.6, baseline: [0.9, 0.9], reasons: [] }).severity).toBe('major') + expect(analyzeCacheRegression({ current: 0.81, baseline: [0.9, 0.92], reasons: [] }).severity).toBe('minor') + }) + + it('prefers the most decisive reason when several are present', () => { + const report = analyzeCacheRegression({ + current: 0.1, + baseline: [0.9, 0.9], + reasons: ['provider_cache_miss', 'tool_catalog_changed', 'cache_ttl_unknown'] + }) + expect(report.primaryReason).toBe('tool_catalog_changed') + }) + + it('stays silent until enough baseline samples exist', () => { + const report = analyzeCacheRegression({ + current: 0.0, + baseline: [0.95], + reasons: ['stable_prefix_changed'], + minBaselineSamples: 2 + }) + expect(report.severity).toBe('none') + expect(report.explanation).toBeNull() + }) + + it('ignores non-rate samples and null current values', () => { + expect(analyzeCacheRegression({ current: null, baseline: [0.9, 0.9], reasons: [] }).severity).toBe('none') + const report = analyzeCacheRegression({ + current: 0.1, + baseline: [null, 0.9, Number.NaN as unknown as number, 0.9], + reasons: ['stable_prefix_changed'] + }) + expect(report.severity).toBe('cliff') + expect(report.baselineHitRate).toBeCloseTo(0.9, 5) + }) + + it('uses a median baseline so a single cold-start outlier does not fake a regression', () => { + // [0.9, 0.0(cold start), 0.9] → median 0.9; current 0.85 is a 5pp dip, not a drop. + const report = analyzeCacheRegression({ current: 0.85, baseline: [0.9, 0.0, 0.9], reasons: [] }) + expect(report.baselineHitRate).toBeCloseTo(0.9, 5) + expect(report.severity).toBe('none') + }) +}) +describe('explainCacheRegression', () => { + it('formats a percentage-point drop with the attributed cause', () => { + const text = explainCacheRegression({ baseline: 0.92, current: 0.08, primaryReason: 'tool_catalog_changed' }) + expect(text).toBe('Cache hit rate dropped from 92% to 8% (-84pp). Likely cause: the tool catalog changed (MCP/Skill tools), which invalidated the cached prefix.') + }) + + it('omits the cause clause when no reason is known', () => { + const text = explainCacheRegression({ baseline: 0.8, current: 0.5, primaryReason: null }) + expect(text).toBe('Cache hit rate dropped from 80% to 50% (-30pp).') + }) +}) diff --git a/kun/src/cache/cache-regression.ts b/kun/src/cache/cache-regression.ts new file mode 100644 index 000000000..6c2fe8f2c --- /dev/null +++ b/kun/src/cache/cache-regression.ts @@ -0,0 +1,154 @@ +import type { CacheMissReason } from './cache-diagnostics.js' + +/** + * Trend-based cache-hit-rate regression analysis. + * + * {@link diagnoseCacheUsage} explains the per-turn miss reasons, but it does + * not know whether the *rate itself dropped* relative to how the thread had + * been performing. This module compares the current cacheable hit rate against + * a rolling baseline of recent turns, classifies the severity of any drop, and + * attributes it to the most likely cause so the GUI can say, for example: + * "Cache hit rate dropped from 92% to 8% (-84pp). Likely cause: the stable + * system prefix changed." + */ + +export type CacheRegressionSeverity = 'none' | 'minor' | 'major' | 'cliff' + +export type CacheRegressionReport = { + severity: CacheRegressionSeverity + baselineHitRate: number | null + currentHitRate: number | null + /** Drop in percentage points (baseline - current) * 100, rounded to 1dp. */ + dropPercentagePoints: number | null + /** The miss reason most likely responsible for the drop, when known. */ + primaryReason: CacheMissReason | null + /** Human-readable explanation, only set when severity is not 'none'. */ + explanation: string | null +} + +// Absolute drop thresholds in hit-rate fraction (0..1). +const MINOR_DROP = 0.08 +const MAJOR_DROP = 0.2 +const CLIFF_DROP = 0.4 + +/** + * Reasons ordered by how decisively they explain a cache regression. The first + * present reason wins, so a prefix change (which invalidates the whole cached + * prefix) is reported ahead of a generic provider miss. + */ +const REASON_PRIORITY: readonly CacheMissReason[] = [ + 'stable_prefix_changed', + 'tool_catalog_changed', + 'skills_changed', + 'model_changed', + 'provider_changed', + 'endpoint_changed', + 'provider_cache_miss', + 'cache_ttl_unknown', + 'cold_request', + 'provider_metrics_unavailable' +] + +const REASON_TEXT: Record = { + cold_request: 'this was the first request in the thread, so there was no warm cache yet', + model_changed: 'the model changed, which starts a new provider cache', + provider_changed: 'the provider changed, which starts a new provider cache', + endpoint_changed: 'the endpoint format changed, which starts a new provider cache', + stable_prefix_changed: 'the stable system prefix changed and invalidated the cached prefix', + tool_catalog_changed: 'the tool catalog changed (MCP/Skill tools), which invalidated the cached prefix', + skills_changed: 'the active Skill set changed, which invalidated the cached prefix', + cache_ttl_unknown: 'the provider cache TTL likely expired before this turn', + provider_cache_miss: 'the provider reported a full cache miss for this turn', + provider_metrics_unavailable: 'the provider did not report cache metrics, so the cause cannot be confirmed' +} + +export function analyzeCacheRegression(input: { + current: number | null + baseline: readonly (number | null)[] + reasons?: readonly CacheMissReason[] + /** Minimum number of usable baseline samples before reporting a drop. */ + minBaselineSamples?: number +}): CacheRegressionReport { + const minSamples = Math.max(1, input.minBaselineSamples ?? 1) + const samples = input.baseline.filter((value): value is number => isRate(value)) + const current = isRate(input.current) ? input.current : null + // Median is robust to a single cold-start zero or one anomalous spike, so a + // lone outlier in the window cannot drag the baseline and fake a regression. + const baseline = samples.length > 0 ? median(samples) : null + const primaryReason = pickPrimaryReason(input.reasons ?? []) + + if (current === null || baseline === null || samples.length < minSamples) { + return { + severity: 'none', + baselineHitRate: baseline, + currentHitRate: current, + dropPercentagePoints: null, + primaryReason, + explanation: null + } + } + + const drop = baseline - current + const severity = classifyDrop(drop) + const dropPercentagePoints = round1(drop * 100) + if (severity === 'none') { + return { severity, baselineHitRate: baseline, currentHitRate: current, dropPercentagePoints, primaryReason, explanation: null } + } + return { + severity, + baselineHitRate: baseline, + currentHitRate: current, + dropPercentagePoints, + primaryReason, + explanation: explainCacheRegression({ baseline, current, primaryReason }) + } +} + +export function explainCacheRegression(input: { + baseline: number + current: number + primaryReason: CacheMissReason | null +}): string { + const drop = round1((input.baseline - input.current) * 100) + const head = `Cache hit rate dropped from ${pct(input.baseline)} to ${pct(input.current)} (-${drop}pp).` + const cause = input.primaryReason ? ` Likely cause: ${REASON_TEXT[input.primaryReason]}.` : '' + return `${head}${cause}` +} + +function classifyDrop(drop: number): CacheRegressionSeverity { + if (drop >= CLIFF_DROP) return 'cliff' + if (drop >= MAJOR_DROP) return 'major' + if (drop >= MINOR_DROP) return 'minor' + return 'none' +} + +function pickPrimaryReason(reasons: readonly CacheMissReason[]): CacheMissReason | null { + for (const reason of REASON_PRIORITY) { + if (reasons.includes(reason)) return reason + } + return null +} + +function isRate(value: unknown): value is number { + return typeof value === 'number' && Number.isFinite(value) && value >= 0 && value <= 1 +} + +/** Ordinal severity for cooldown comparisons; higher means a worse drop. */ +export function cacheRegressionSeverityRank(severity: CacheRegressionSeverity): number { + return { none: 0, minor: 1, major: 2, cliff: 3 }[severity] +} + +/** Median of the sample window — robust to a single cold-start or outlier. */ +function median(values: readonly number[]): number { + const sorted = [...values].sort((a, b) => a - b) + const mid = Math.floor(sorted.length / 2) + return sorted.length % 2 === 0 ? (sorted[mid - 1] + sorted[mid]) / 2 : sorted[mid] +} + +function round1(value: number): number { + return Math.round(value * 10) / 10 +} + +function pct(rate: number): string { + return `${Math.round(rate * 100)}%` +} diff --git a/kun/src/cli/agent-cli.ts b/kun/src/cli/agent-cli.ts index 54b0b11d8..85829b435 100644 --- a/kun/src/cli/agent-cli.ts +++ b/kun/src/cli/agent-cli.ts @@ -43,6 +43,7 @@ Common options: --model Model id --approval-policy

on-request | untrusted | never | auto | suggest --json Emit machine-readable JSON where supported + --jsonl Stream one machine-readable event per line for kun run Exec options: --list-tools Print available tools @@ -110,6 +111,11 @@ export async function runAgentCommand( async function runOneShot(argv: readonly string[], io: CliIo): Promise { const parsed = parseSharedOptions(argv, io) if (!parsed.ok) return writeParseError(parsed, io, 'kun run') + const jsonl = hasFlag(argv, 'jsonl') + if (parsed.json && jsonl) { + io.stderr.write('kun run: --json and --jsonl are mutually exclusive\n') + return ServeExitCode.usage + } const prompt = stringFlag(argv, ['prompt', 'p']) ?? positionals(argv).join(' ').trim() if (!prompt) { io.stderr.write('kun run: missing prompt\n') @@ -126,13 +132,16 @@ async function runOneShot(argv: readonly string[], io: CliIo): Promise { approvalPolicy: parsed.options.approvalPolicy, sandboxMode: parsed.options.sandboxMode }) + if (jsonl) writeJsonLine(io.stdout, { type: 'run_started', threadId: thread.id }) const turn = await runtime.turnService.startTurn({ threadId: thread.id, request: { prompt, model: parsed.options.model, mode: 'agent' } }) let streamed = false - const unsubscribe = parsed.json ? undefined : runtime.eventBus.subscribe(thread.id, (event) => { - if (event.kind === 'assistant_text_delta' && event.item.kind === 'assistant_text') { + const unsubscribe = runtime.eventBus.subscribe(thread.id, (event) => { + if (jsonl) { + writeJsonLine(io.stdout, { type: 'runtime_event', event }) + } else if (!parsed.json && event.kind === 'assistant_text_delta' && event.item.kind === 'assistant_text') { streamed = true io.stdout.write(event.item.text) } @@ -140,7 +149,9 @@ async function runOneShot(argv: readonly string[], io: CliIo): Promise { const status = await runtime.runTurn(thread.id, turn.turnId) unsubscribe?.() const items = await runtime.sessionStore.loadItems(thread.id) - if (parsed.json) { + if (jsonl) { + writeJsonLine(io.stdout, { type: 'run_finished', threadId: thread.id, turnId: turn.turnId, status }) + } else if (parsed.json) { io.stdout.write(JSON.stringify({ threadId: thread.id, turnId: turn.turnId, status, items }) + '\n') } else { if (!streamed) { @@ -158,6 +169,10 @@ async function runOneShot(argv: readonly string[], io: CliIo): Promise { } } +function writeJsonLine(output: WritableLike, value: unknown): void { + output.write(`${JSON.stringify(value)}\n`) +} + async function runChat(argv: readonly string[], io: CliIo): Promise { const parsed = parseSharedOptions(argv, io) if (!parsed.ok) return writeParseError(parsed, io, 'kun chat') diff --git a/kun/src/contracts/attachments.ts b/kun/src/contracts/attachments.ts index b253a0029..d4fe984a4 100644 --- a/kun/src/contracts/attachments.ts +++ b/kun/src/contracts/attachments.ts @@ -13,11 +13,15 @@ export type AttachmentTextFallback = z.infer export const AttachmentMetadata = z.object({ id: z.string().min(1), name: z.string().min(1), + kind: z.enum(['image', 'document']).default('image'), mimeType: z.string().min(1), byteSize: z.number().int().nonnegative(), hash: z.string().min(1), width: z.number().int().positive().optional(), height: z.number().int().positive().optional(), + documentText: z.string().optional(), + pageCount: z.number().int().positive().optional(), + truncated: z.boolean().optional(), localFilePath: z.string().min(1).optional(), textFallback: AttachmentTextFallback.optional(), threadIds: z.array(z.string().min(1)).default([]), @@ -31,6 +35,8 @@ export const AttachmentUploadRequest = z.object({ name: z.string().min(1), mimeType: z.string().min(1).optional(), dataBase64: z.string().min(1), + documentText: z.string().optional(), + pageCount: z.number().int().positive().optional(), localFilePath: z.string().min(1).optional(), textFallback: AttachmentTextFallback.optional(), threadId: z.string().min(1).optional(), diff --git a/kun/src/contracts/capabilities.ts b/kun/src/contracts/capabilities.ts index 95e9f3859..03024c358 100644 --- a/kun/src/contracts/capabilities.ts +++ b/kun/src/contracts/capabilities.ts @@ -83,6 +83,19 @@ export type McpTrustScope = z.infer export const McpToolDiscoveryMode = z.enum(['direct', 'search', 'auto']) export type McpToolDiscoveryMode = z.infer +export const McpOAuthConfig = z + .object({ + enabled: z.boolean().default(true), + clientName: z.string().min(1).optional(), + clientId: z.string().min(1).optional(), + clientSecret: z.string().min(1).optional(), + scopes: z.array(z.string().min(1)).default([]), + redirectPort: z.number().int().min(1024).max(65535).optional(), + callbackTimeoutMs: z.number().int().positive().default(120_000) + }) + .strict() +export type McpOAuthConfig = z.infer + export const McpSearchConfig = z .object({ enabled: z.boolean().default(false), @@ -124,6 +137,7 @@ export const McpServerConfig = z // Visibility scope: empty means globally visible; otherwise the server is // advertised only when ToolHostContext.workspace is under one of these roots. workspaceRoots: z.array(z.string().min(1)).default([]), + oauth: McpOAuthConfig.optional(), trustScope: McpTrustScope.default('workspace'), trustedWorkspaceRoots: z.array(z.string().min(1)).default([]), timeoutMs: z.number().int().positive().default(30_000) @@ -219,7 +233,7 @@ export type SubagentMode = z.infer * to side-effect-free investigation tools — no bash/edit/write, and no * nested `delegate_task`. */ -export const SUBAGENT_READ_ONLY_TOOL_NAMES = ['read', 'grep', 'find', 'ls'] as const +export const SUBAGENT_READ_ONLY_TOOL_NAMES = ['read', 'grep', 'find', 'ls', 'repo_map'] as const export const SubagentProfileConfig = z .object({ @@ -302,11 +316,23 @@ export type SubagentsCapabilityConfig = z.output +export const MemorySourceKind = z.enum(['user', 'tool', 'inference', 'file', 'web']) +export type MemorySourceKind = z.infer + +export const MemoryProvenance = z.object({ + kind: MemorySourceKind, + turnId: z.string().optional(), + file: z.string().optional(), + origin: z.string().optional() +}).strict() +export type MemoryProvenance = z.infer + export const MemoryRecord = z.object({ id: z.string().min(1), content: z.string().min(1), @@ -11,10 +22,15 @@ export const MemoryRecord = z.object({ project: z.string().optional(), sourceThreadId: z.string().optional(), sourceTurnId: z.string().optional(), + provenance: MemoryProvenance.optional(), tags: z.array(z.string()).default([]), confidence: z.number().min(0).max(1).default(1), createdAt: z.string(), updatedAt: z.string(), + expiresAt: z.string().datetime().optional(), + supersedes: z.string().optional(), + supersededAt: z.string().optional(), + correctedFrom: z.string().optional(), disabledAt: z.string().optional(), deletedAt: z.string().optional() }).strict() @@ -27,8 +43,11 @@ export const MemoryCreateRequest = z.object({ project: z.string().optional(), sourceThreadId: z.string().optional(), sourceTurnId: z.string().optional(), + provenance: MemoryProvenance.optional(), + ttlMs: z.number().int().positive().optional(), + supersedes: z.string().min(1).optional(), tags: z.array(z.string()).default([]), - confidence: z.number().min(0).max(1).default(1) + confidence: z.number().min(0).max(1).optional() }).strict() export type MemoryCreateRequest = z.input @@ -36,6 +55,7 @@ export const MemoryUpdateRequest = z.object({ content: z.string().min(1).optional(), tags: z.array(z.string()).optional(), confidence: z.number().min(0).max(1).optional(), + expiresAt: z.string().datetime().nullable().optional(), disabled: z.boolean().optional() }).strict() export type MemoryUpdateRequest = z.input diff --git a/kun/src/delegation/child-agent-executor.ts b/kun/src/delegation/child-agent-executor.ts index 95e9c3ad5..87584994b 100644 --- a/kun/src/delegation/child-agent-executor.ts +++ b/kun/src/delegation/child-agent-executor.ts @@ -16,6 +16,7 @@ import { InflightTracker } from '../loop/inflight-tracker.js' import { SteeringQueue } from '../loop/steering-queue.js' import type { TokenEconomyConfig } from '../loop/token-economy.js' import type { MemoryStore } from '../memory/memory-store.js' +import type { ArtifactStore } from '../artifacts/artifact-store.js' import type { ModelClient } from '../ports/model-client.js' import { RandomIdGenerator } from '../ports/id-generator.js' import type { SessionStore } from '../ports/session-store.js' @@ -43,6 +44,7 @@ export type ChildAgentExecutorOptions = { modelCapabilities?: (model: string) => ModelCapabilityMetadata skillRuntime?: SkillRuntime memoryStore?: MemoryStore + artifactStore?: ArtifactStore /** * Persistence wiring. When the main runtime's stores + event recorder are * supplied, the child runs as a persisted `relation: 'side'` thread on the @@ -159,6 +161,7 @@ export function createChildAgentExecutor(options: ChildAgentExecutorOptions): Ch ...(options.modelCapabilities ? { modelCapabilities: options.modelCapabilities } : {}), ...(options.skillRuntime ? { skillRuntime: options.skillRuntime } : {}), ...(options.memoryStore ? { memoryStore: options.memoryStore } : {}), + ...(options.artifactStore ? { artifactStore: options.artifactStore } : {}), ...(options.contextCompaction ? { contextCompaction: options.contextCompaction } : {}), ...(options.tokenEconomy ? { tokenEconomy: options.tokenEconomy } : {}), ...(options.runtime?.toolStorm ? { toolStorm: options.runtime.toolStorm } : {}), @@ -189,9 +192,12 @@ export function createChildAgentExecutor(options: ChildAgentExecutorOptions): Ch }) // A profile preamble rides in the prompt body (not the system prompt) so // the cached stable prefix stays byte-identical to the main agent's. - const prompt = input.promptPreamble?.trim() + const promptBase = input.promptPreamble?.trim() ? `${input.promptPreamble.trim()}\n\n${input.prompt}` : input.prompt + const prompt = input.returnFormat === 'evidence' + ? `${promptBase}\n\nReturn a concise evidence-based conclusion. Inspect the task with tools so the parent can verify the result.` + : promptBase const started = await turns.startTurn({ threadId: thread.id, request: { @@ -227,11 +233,15 @@ export function createChildAgentExecutor(options: ChildAgentExecutorOptions): Ch const toolInvocations = items.filter( (item) => item.turnId === started.turnId && item.kind === 'tool_call' ).length + const evidence = input.returnFormat === 'evidence' + ? childToolEvidence(items, started.turnId) + : undefined if (status !== 'completed') { throw new Error(summary || `child agent ${status}`) } return { summary, + ...(evidence ? { evidence } : {}), usage: usage.forThread(thread.id), toolInvocations, // The child loop was constructed with the main agent's immutable @@ -242,6 +252,30 @@ export function createChildAgentExecutor(options: ChildAgentExecutorOptions): Ch } } +function childToolEvidence(items: readonly TurnItem[], turnId: string): string[] { + const results = new Map(items + .filter((item): item is Extract => + item.turnId === turnId && item.kind === 'tool_result') + .map((item) => [item.callId, item])) + return items + .filter((item): item is Extract => + item.turnId === turnId && item.kind === 'tool_call') + .slice(0, 32) + .map((item) => { + const result = results.get(item.callId) + const target = toolEvidenceTarget(item.arguments) + return `${item.toolName}${target ? ` ${target}` : ''}: ${result?.isError ? 'failed' : 'completed'}` + }) +} + +function toolEvidenceTarget(args: Record): string { + for (const key of ['path', 'filePath', 'file_path', 'query', 'command']) { + const value = args[key] + if (typeof value === 'string' && value.trim()) return value.trim().slice(0, 300) + } + return '' +} + function childThreadTitle(childId: string, label?: string, profile?: string): string { const suffix = label?.trim() || profile?.trim() || childId return `Child agent: ${suffix}` diff --git a/kun/src/delegation/delegation-runtime.ts b/kun/src/delegation/delegation-runtime.ts index 3b2493804..bb6bd5532 100644 --- a/kun/src/delegation/delegation-runtime.ts +++ b/kun/src/delegation/delegation-runtime.ts @@ -24,6 +24,9 @@ const ChildRunUsage = z.object({ tokenEconomySavingsCny: z.number().nonnegative().optional() }) +const ChildReturnFormat = z.enum(['summary', 'evidence']) +export type ChildReturnFormat = z.infer + export const ChildRunRecord = z.object({ id: z.string().min(1), parentThreadId: z.string().min(1), @@ -40,6 +43,11 @@ export const ChildRunRecord = z.object({ toolPolicy: SubagentToolPolicy.optional(), status: z.enum(['queued', 'running', 'completed', 'failed', 'aborted']), summary: z.string().optional(), + evidence: z.array(z.string().min(1).max(2_000)).max(32).optional(), + tokenBudget: z.number().int().positive().optional(), + timeBudgetMs: z.number().int().positive().optional(), + returnFormat: ChildReturnFormat.default('summary'), + budgetExceeded: z.enum(['token', 'time']).optional(), error: z.string().optional(), usage: ChildRunUsage.default({ promptTokens: 0, completionTokens: 0, totalTokens: 0 }), /** True when the child reused the main agent's cached stable prefix. */ @@ -82,6 +90,9 @@ export type ChildRunExecutor = (input: { promptPreamble?: string /** Reasoning depth for this profile's child model requests (default 'off'). */ reasoningEffort?: string + tokenBudget?: number + timeBudgetMs?: number + returnFormat?: ChildReturnFormat signal: AbortSignal }) => Promise<{ summary: string @@ -89,6 +100,7 @@ export type ChildRunExecutor = (input: { toolInvocations?: number prefixReused?: boolean inheritedHistoryItems?: number + evidence?: string[] }> export type ChildRunAggregate = { @@ -174,6 +186,9 @@ export class DelegationRuntime { model?: string providerId?: string profile?: string + tokenBudget?: number + timeBudgetMs?: number + returnFormat?: ChildReturnFormat /** * When true, runChild returns the queued ChildRunRecord immediately and * continues execution in the background. The detached run gets its own @@ -219,6 +234,9 @@ export class DelegationRuntime { const resolvedBlockedSkills = profile?.blockedSkills const promptPreamble = profile?.promptPreamble const resolvedReasoningEffort = profile?.reasoningEffort + const tokenBudget = positiveInteger(input.tokenBudget) + const timeBudgetMs = positiveInteger(input.timeBudgetMs) + const returnFormat = input.returnFormat ?? 'summary' // Reserve against the per-thread budget before persisting anything. await this.ensureSeeded(input.parentThreadId) @@ -239,6 +257,9 @@ export class DelegationRuntime { providerId: resolvedProviderId, profile: profileName, toolPolicy, + tokenBudget, + timeBudgetMs, + returnFormat, status: 'queued', createdAt: queuedAt, updatedAt: queuedAt @@ -271,6 +292,9 @@ export class DelegationRuntime { resolvedBlockedSkills, promptPreamble, resolvedReasoningEffort, + tokenBudget, + timeBudgetMs, + returnFormat, workspace: input.workspace, label: input.label, parentThreadId: input.parentThreadId, @@ -303,7 +327,7 @@ export class DelegationRuntime { await this.recordChildEvent(record) try { const executor: ChildRunExecutor = this.options.executor ?? defaultExecutor - const result = await executor({ + const result = await executeWithinTimeBudget(input.signal, timeBudgetMs, (signal) => executor({ childId: id, parentThreadId: input.parentThreadId, parentTurnId: input.parentTurnId, @@ -321,17 +345,31 @@ export class DelegationRuntime { toolPolicy, ...(promptPreamble ? { promptPreamble } : {}), ...(resolvedReasoningEffort ? { reasoningEffort: resolvedReasoningEffort } : {}), - signal: input.signal - }) + ...(tokenBudget ? { tokenBudget } : {}), + ...(timeBudgetMs ? { timeBudgetMs } : {}), + returnFormat, + signal + })) const finishedAt = this.now() + const usage = result.usage ?? record.usage + const contractError = childContractError({ tokenBudget, returnFormat }, result.evidence, usage) record = ChildRunRecord.parse({ ...record, - status: 'completed', + status: contractError ? 'failed' : 'completed', summary: result.summary, - usage: result.usage ?? record.usage, + evidence: result.evidence, + usage, toolInvocations: result.toolInvocations, prefixReused: result.prefixReused, inheritedHistoryItems: result.inheritedHistoryItems, + ...(contractError + ? { + error: contractError, + ...(usage.totalTokens > (tokenBudget ?? Number.POSITIVE_INFINITY) + ? { budgetExceeded: 'token' as const } + : {}) + } + : {}), durationMs: elapsedMs(startedAt, finishedAt), updatedAt: finishedAt }) @@ -345,6 +383,7 @@ export class DelegationRuntime { ...record, status: input.signal.aborted ? 'aborted' : 'failed', error: errorMessage(error), + ...(error instanceof ChildTimeBudgetExceededError ? { budgetExceeded: 'time' } : {}), durationMs: elapsedMs(startedAt, finishedAt), updatedAt: finishedAt }) @@ -377,6 +416,9 @@ export class DelegationRuntime { resolvedBlockedSkills: string[] | undefined promptPreamble: string | undefined resolvedReasoningEffort: string | undefined + tokenBudget: number | undefined + timeBudgetMs: number | undefined + returnFormat: ChildReturnFormat workspace: string | undefined label: string | undefined parentThreadId: string @@ -406,7 +448,7 @@ export class DelegationRuntime { await this.recordChildEvent(record) try { const executor: ChildRunExecutor = this.options.executor ?? defaultExecutor - const result = await executor({ + const result = await executeWithinTimeBudget(args.signal, args.timeBudgetMs, (signal) => executor({ childId: record.id, parentThreadId: args.parentThreadId, parentTurnId: args.parentTurnId, @@ -424,17 +466,35 @@ export class DelegationRuntime { toolPolicy: args.toolPolicy, ...(args.promptPreamble ? { promptPreamble: args.promptPreamble } : {}), ...(args.resolvedReasoningEffort ? { reasoningEffort: args.resolvedReasoningEffort } : {}), - signal: args.signal - }) + ...(args.tokenBudget ? { tokenBudget: args.tokenBudget } : {}), + ...(args.timeBudgetMs ? { timeBudgetMs: args.timeBudgetMs } : {}), + returnFormat: args.returnFormat, + signal + })) const finishedAt = this.now() + const usage = result.usage ?? record.usage + const contractError = childContractError( + { tokenBudget: args.tokenBudget, returnFormat: args.returnFormat }, + result.evidence, + usage + ) record = ChildRunRecord.parse({ ...record, - status: 'completed', + status: contractError ? 'failed' : 'completed', summary: result.summary, - usage: result.usage ?? record.usage, + evidence: result.evidence, + usage, toolInvocations: result.toolInvocations, prefixReused: result.prefixReused, inheritedHistoryItems: result.inheritedHistoryItems, + ...(contractError + ? { + error: contractError, + ...(usage.totalTokens > (args.tokenBudget ?? Number.POSITIVE_INFINITY) + ? { budgetExceeded: 'token' as const } + : {}) + } + : {}), durationMs: elapsedMs(startedAt, finishedAt), updatedAt: finishedAt }) @@ -448,6 +508,7 @@ export class DelegationRuntime { ...record, status: args.signal.aborted ? 'aborted' : 'failed', error: errorMessage(error), + ...(error instanceof ChildTimeBudgetExceededError ? { budgetExceeded: 'time' } : {}), durationMs: elapsedMs(startedAt, finishedAt), updatedAt: finishedAt }) @@ -635,7 +696,6 @@ export class DelegationRuntime { } private recordExternalUsage(record: ChildRunRecord): void { - if (record.status !== 'completed') return const usage = toUsageSnapshot(record.usage) if (usage.totalTokens <= 0 && usage.costUsd === undefined && usage.costCny === undefined) return this.options.recordExternalUsage?.(record.parentThreadId, usage) @@ -706,6 +766,67 @@ export function aggregateChildRuns(records: readonly ChildRunRecord[]): ChildRun ) } +class ChildTimeBudgetExceededError extends Error { + constructor(readonly timeBudgetMs: number) { + super(`child time budget exhausted after ${timeBudgetMs}ms`) + this.name = 'ChildTimeBudgetExceededError' + } +} + +async function executeWithinTimeBudget( + parentSignal: AbortSignal, + timeBudgetMs: number | undefined, + execute: (signal: AbortSignal) => Promise +): Promise { + if (parentSignal.aborted) throw new Error('child run aborted') + if (!timeBudgetMs) return execute(parentSignal) + + const controller = new AbortController() + let timer: ReturnType | undefined + let rejectCancellation: ((error: Error) => void) | undefined + const cancellation = new Promise((_resolve, reject) => { + rejectCancellation = reject + }) + const onParentAbort = (): void => { + controller.abort(parentSignal.reason) + rejectCancellation?.(new Error('child run aborted')) + } + parentSignal.addEventListener('abort', onParentAbort, { once: true }) + timer = setTimeout(() => { + const error = new ChildTimeBudgetExceededError(timeBudgetMs) + controller.abort(error) + rejectCancellation?.(error) + }, timeBudgetMs) + timer.unref?.() + + try { + return await Promise.race([execute(controller.signal), cancellation]) + } finally { + if (timer) clearTimeout(timer) + parentSignal.removeEventListener('abort', onParentAbort) + } +} + +function childContractError( + contract: { tokenBudget: number | undefined; returnFormat: ChildReturnFormat }, + evidence: string[] | undefined, + usage: ChildRunRecord['usage'] +): string | undefined { + if (contract.tokenBudget !== undefined && usage.totalTokens > contract.tokenBudget) { + return `child token budget exhausted (${usage.totalTokens} > ${contract.tokenBudget})` + } + if (contract.returnFormat === 'evidence' && !evidence?.some((item) => item.trim().length > 0)) { + return 'child contract requires evidence but none was returned' + } + return undefined +} + +function positiveInteger(value: number | undefined): number | undefined { + if (value === undefined) return undefined + if (!Number.isInteger(value) || value <= 0) throw new Error('child budgets must be positive integers') + return value +} + const defaultExecutor: ChildRunExecutor = async (input) => { return { summary: `Child result: ${input.prompt}` } } diff --git a/kun/src/loop/agent-loop.ts b/kun/src/loop/agent-loop.ts index 40ed79672..550f96f14 100644 --- a/kun/src/loop/agent-loop.ts +++ b/kun/src/loop/agent-loop.ts @@ -61,8 +61,9 @@ import type { ThreadGoal, ThreadTodoList } from '../contracts/threads.js' import { modelCapabilitiesForModel, type ContextCompactionConfig } from './model-context-profile.js' import type { SkillRuntime } from '../skills/skill-runtime.js' import type { AttachmentContent, AttachmentStore } from '../attachments/attachment-store.js' -import type { ModelInputAttachment, ModelTextAttachmentFallback } from '../ports/model-client.js' +import type { ModelDocumentAttachment, ModelInputAttachment, ModelTextAttachmentFallback } from '../ports/model-client.js' import type { MemoryStore } from '../memory/memory-store.js' +import type { ArtifactStore } from '../artifacts/artifact-store.js' import { hasHooksForPhase, runObserverHooks, @@ -91,6 +92,7 @@ import { GET_GOAL_TOOL_NAME, UPDATE_GOAL_TOOL_NAME } from '../adapters/tool/goal import { TODO_LIST_TOOL_NAME, TODO_WRITE_TOOL_NAME } from '../adapters/tool/todo-tools.js' import { shellRuntimeInstruction } from '../adapters/tool/builtin-tool-utils.js' import { VERIFY_CHANGES_TOOL_NAME } from '../adapters/tool/builtin-verify-tool.js' +import { buildToolPreferenceInstruction } from '../prompt/kun-system-prompt.js' import { GoalResumeCoordinator, DEFAULT_MAX_GOAL_RESUME_NO_PROGRESS_ATTEMPTS, @@ -672,6 +674,7 @@ export type AgentLoopOptions = { skillRuntime?: SkillRuntime attachmentStore?: AttachmentStore memoryStore?: MemoryStore + artifactStore?: ArtifactStore /** Kun runtime data root for sandbox-safe background shell output reads. */ runtimeDataDir?: string tokenEconomy?: TokenEconomyConfig @@ -1534,6 +1537,7 @@ export class AgentLoop { nowIso: this.opts.nowIso() }) : null + const toolPreferenceInstruction = buildToolPreferenceInstruction(tools) const contextInstructions = [ ...(runtimeContextInstruction ? [runtimeContextInstruction] : []), ...(activeGoalInstruction ? [activeGoalInstruction] : []), @@ -1554,6 +1558,7 @@ export class AgentLoop { ...(skillResolution.catalogInstruction ? [skillResolution.catalogInstruction] : []), ...skillResolution.instructions, ...(userInputDisabled ? [userInputUnavailableInstruction()] : []), + ...(toolPreferenceInstruction ? [toolPreferenceInstruction] : []), ...(effectiveToolSpecs.some((tool) => tool.name === 'bash') ? [shellRuntimeInstruction()] : []), ...(suggestVerification ? [verificationSuggestionInstruction()] : []), ...(toolCatalogDriftMessage ? [toolCatalogDriftMessage] : []) @@ -1583,6 +1588,7 @@ export class AgentLoop { history: capToolResultImages(history, MAX_FORWARDED_TOOL_IMAGES), ...(attachments.imageAttachments.length ? { attachments: attachments.imageAttachments } : {}), ...(attachments.textFallbacks.length ? { attachmentTextFallbacks: attachments.textFallbacks } : {}), + ...(attachments.documents.length ? { attachmentDocuments: attachments.documents } : {}), tools: effectiveToolSpecs, ...(requiredToolName ? { requiredToolName } : {}), ...(modelRoute.reasoningEffort ? { reasoningEffort: modelRoute.reasoningEffort } : {}), @@ -1664,6 +1670,7 @@ export class AgentLoop { attachmentIds: turn?.attachmentIds ?? [], imageAttachments: attachments.imageAttachments, textFallbacks: attachments.textFallbacks, + documents: attachments.documents, modelCapabilities }) }) @@ -2225,6 +2232,7 @@ export class AgentLoop { approvalPolicy: input.approvalPolicy, sandboxMode: input.sandboxMode, ...(this.opts.runtimeDataDir ? { runtimeDataDir: this.opts.runtimeDataDir } : {}), + ...(this.opts.artifactStore ? { artifactStore: this.opts.artifactStore } : {}), abortSignal: input.signal, awaitApproval: async (approval) => { await this.opts.events.record({ @@ -2968,8 +2976,8 @@ export class AgentLoop { threadId: string workspace: string modelCapabilities: ModelCapabilityMetadata - }): Promise<{ imageAttachments: ModelInputAttachment[]; textFallbacks: ModelTextAttachmentFallback[] }> { - if (input.attachmentIds.length === 0) return { imageAttachments: [], textFallbacks: [] } + }): Promise<{ imageAttachments: ModelInputAttachment[]; textFallbacks: ModelTextAttachmentFallback[]; documents: ModelDocumentAttachment[] }> { + if (input.attachmentIds.length === 0) return { imageAttachments: [], textFallbacks: [], documents: [] } if (!this.opts.attachmentStore) { throw new Error('attachment store is unavailable') } @@ -2977,11 +2985,30 @@ export class AgentLoop { const textFallbackPolicy = this.opts.attachmentStore.textFallbackPolicy() const imageAttachments: ModelInputAttachment[] = [] const textFallbacks: ModelTextAttachmentFallback[] = [] + const documents: ModelDocumentAttachment[] = [] + let remainingDocumentChars = 400_000 for (const id of input.attachmentIds) { const attachment = await this.opts.attachmentStore.resolveContent(id, { threadId: input.threadId, workspace: input.workspace }) + if (attachment.kind === 'document') { + const fullText = attachment.documentText ?? '' + const text = fullText.slice(0, Math.max(0, remainingDocumentChars)) + remainingDocumentChars -= text.length + documents.push({ + id: attachment.id, + name: attachment.name, + mimeType: attachment.mimeType, + text, + byteSize: attachment.byteSize, + ...(attachment.pageCount ? { pageCount: attachment.pageCount } : {}), + ...(attachment.truncated || text.length < fullText.length ? { truncated: true } : {}), + ...(attachment.localFilePath ? { localFilePath: attachment.localFilePath } : {}) + }) + if (remainingDocumentChars <= 0) break + continue + } if (supportsImageInput) { imageAttachments.push({ id: attachment.id, @@ -2999,7 +3026,7 @@ export class AgentLoop { textFallbackPolicy.textFallbackMaxBase64Bytes )) } - return { imageAttachments, textFallbacks } + return { imageAttachments, textFallbacks, documents } } private async retrieveMemories(input: { @@ -3071,12 +3098,15 @@ function attachmentRequestPipelineDetails(input: { attachmentIds: readonly string[] imageAttachments: readonly ModelInputAttachment[] textFallbacks: readonly ModelTextAttachmentFallback[] + documents?: readonly ModelDocumentAttachment[] modelCapabilities: ModelCapabilityMetadata }): Record { + const documents = input.documents ?? [] if ( input.attachmentIds.length === 0 && input.imageAttachments.length === 0 && - input.textFallbacks.length === 0 + input.textFallbacks.length === 0 && + documents.length === 0 ) { return {} } @@ -3095,7 +3125,10 @@ function attachmentRequestPipelineDetails(input: { (total, attachment) => total + Buffer.byteLength(attachment.dataBase64, 'utf8'), 0 ), - textFallbackMimeTypes: [...new Set(input.textFallbacks.map((attachment) => attachment.mimeType))] + textFallbackMimeTypes: [...new Set(input.textFallbacks.map((attachment) => attachment.mimeType))], + documentCount: documents.length, + documentTextChars: documents.reduce((total, document) => total + document.text.length, 0), + documentMimeTypes: [...new Set(documents.map((document) => document.mimeType))] } } diff --git a/kun/src/memory/memory-store.ts b/kun/src/memory/memory-store.ts index 1134ab92b..c1542270a 100644 --- a/kun/src/memory/memory-store.ts +++ b/kun/src/memory/memory-store.ts @@ -5,10 +5,13 @@ import { atomicWriteFile } from '../adapters/file/atomic-write.js' import { MemoryDiagnostics, MemoryRecord, + type MemoryProvenance, type MemoryCreateRequest, type MemoryUpdateRequest } from '../contracts/memory.js' +const DEFAULT_MEMORY_CONFIDENCE_HALF_LIFE_MS = 180 * 24 * 60 * 60 * 1_000 + export interface MemoryStore { create(input: MemoryCreateRequest): Promise update(id: string, patch: MemoryUpdateRequest, access?: MemoryAccess): Promise @@ -30,6 +33,8 @@ export class FileMemoryStore implements MemoryStore { config: MemoryCapabilityConfig nowIso?: () => string idGenerator?: () => string + confidenceHalfLifeMs?: number + minConfidence?: number } ) {} @@ -39,6 +44,7 @@ export class FileMemoryStore implements MemoryStore { const scope = input.scope ?? 'workspace' const workspace = normalizeScopePath(input.workspace) const project = normalizeScopePath(input.project ?? (scope === 'project' ? input.workspace : undefined)) + const provenance = input.provenance ?? defaultProvenance(input) const parsed = MemoryRecord.parse({ id: this.options.idGenerator?.() ?? `mem_${Date.now().toString(36)}_${Math.random().toString(36).slice(2, 8)}`, content: input.content, @@ -47,11 +53,21 @@ export class FileMemoryStore implements MemoryStore { ...(scope === 'project' && project ? { project } : {}), sourceThreadId: input.sourceThreadId, sourceTurnId: input.sourceTurnId, + provenance, tags: input.tags ?? [], - confidence: input.confidence ?? 1, + confidence: input.confidence ?? defaultConfidence(provenance.kind), createdAt: now, - updatedAt: now + updatedAt: now, + ...(input.ttlMs ? { expiresAt: new Date(Date.parse(now) + input.ttlMs).toISOString() } : {}), + ...(input.supersedes ? { supersedes: input.supersedes } : {}) }) + if (input.supersedes) { + const older = await this.mustGet(input.supersedes, { workspace }) + if (older.scope !== parsed.scope) { + throw new Error('a memory can only supersede another memory in the same scope') + } + await this.write(MemoryRecord.parse({ ...older, supersededAt: now, updatedAt: now })) + } await this.write(parsed) return parsed } @@ -59,11 +75,23 @@ export class FileMemoryStore implements MemoryStore { async update(id: string, patch: MemoryUpdateRequest, access?: MemoryAccess): Promise { const current = await this.mustGet(id, access) const now = this.now() + const corrected = patch.content !== undefined && patch.content !== current.content const next = MemoryRecord.parse({ ...current, ...(patch.content !== undefined ? { content: patch.content } : {}), ...(patch.tags !== undefined ? { tags: patch.tags } : {}), - ...(patch.confidence !== undefined ? { confidence: patch.confidence } : {}), + ...(patch.confidence !== undefined + ? { confidence: patch.confidence } + : corrected + ? { confidence: 1 } + : {}), + ...(corrected + ? { + correctedFrom: current.correctedFrom ?? current.content, + provenance: { ...(current.provenance ?? defaultLegacyProvenance(current)), kind: 'user' } + } + : {}), + ...(patch.expiresAt !== undefined ? { expiresAt: patch.expiresAt ?? undefined } : {}), ...(patch.disabled === true ? { disabledAt: current.disabledAt ?? now } : {}), ...(patch.disabled === false ? { disabledAt: undefined } : {}), updatedAt: now @@ -94,8 +122,14 @@ export class FileMemoryStore implements MemoryStore { async retrieve(input: { query: string; workspace?: string; limit: number }): Promise { if (!this.options.config.enabled) return [] + const nowMs = Date.parse(this.now()) const active = (await this.list({ workspace: input.workspace })) - .filter((record) => !record.disabledAt) + .filter((record) => isMemoryActive( + record, + nowMs, + this.options.minConfidence ?? 0, + this.options.confidenceHalfLifeMs ?? DEFAULT_MEMORY_CONFIDENCE_HALF_LIFE_MS + )) // User-scope memories are persistent identity facts (name, preferences, // account) — small in number, high in value, and frequently queried by // semantic prompts ("who am I?", "what do you know about me") that share @@ -105,7 +139,12 @@ export class FileMemoryStore implements MemoryStore { const userMemories = active.filter((record) => record.scope === 'user') const scored = active .filter((record) => record.scope !== 'user') - .map((record) => ({ record, score: scoreMemory(record, input.query) })) + .map((record) => ({ record, score: scoreMemory( + record, + input.query, + nowMs, + this.options.confidenceHalfLifeMs ?? DEFAULT_MEMORY_CONFIDENCE_HALF_LIFE_MS + ) })) .filter((entry) => entry.score > 0) .sort((a, b) => b.score - a.score || b.record.updatedAt.localeCompare(a.record.updatedAt)) .map((entry) => entry.record) @@ -114,10 +153,16 @@ export class FileMemoryStore implements MemoryStore { async diagnostics(): Promise { const records = await this.readAll() + const nowMs = Date.parse(this.now()) return { enabled: this.options.config.enabled, rootDir: this.options.rootDir, - activeCount: records.filter((record) => !record.deletedAt && !record.disabledAt).length, + activeCount: records.filter((record) => isMemoryActive( + record, + nowMs, + this.options.minConfidence ?? 0, + this.options.confidenceHalfLifeMs ?? DEFAULT_MEMORY_CONFIDENCE_HALF_LIFE_MS + )).length, tombstoneCount: records.filter((record) => Boolean(record.deletedAt)).length, lastInjectedIds: [...this.lastInjectedIds] } @@ -176,7 +221,12 @@ function normalizeScopePath(value: string | undefined): string | undefined { return process.platform === 'win32' ? normalized.toLowerCase() : normalized } -function scoreMemory(record: MemoryRecord, query: string): number { +function scoreMemory( + record: MemoryRecord, + query: string, + nowMs: number, + confidenceHalfLifeMs: number +): number { // Build n-gram fingerprints so matching works for both Latin words and CJK // text. The previous implementation split on `[^a-z0-9_]+`, which treated // every Chinese/Japanese/Korean character as a separator and produced an @@ -190,7 +240,61 @@ function scoreMemory(record: MemoryRecord, query: string): number { } // Normalize by query coverage so long queries do not drown out short ones. const coverage = overlap / queryGrams.size - return (overlap + coverage) * record.confidence + return (overlap + coverage) * effectiveMemoryConfidence(record, nowMs, confidenceHalfLifeMs) +} + +export function isMemoryActive( + record: MemoryRecord, + nowMs: number, + minConfidence = 0, + halfLifeMs = DEFAULT_MEMORY_CONFIDENCE_HALF_LIFE_MS +): boolean { + if (record.deletedAt || record.disabledAt || record.supersededAt) return false + if (record.expiresAt && Date.parse(record.expiresAt) <= nowMs) return false + return effectiveMemoryConfidence( + record, + nowMs, + halfLifeMs + ) >= minConfidence +} + +export function effectiveMemoryConfidence( + record: MemoryRecord, + nowMs: number, + halfLifeMs = DEFAULT_MEMORY_CONFIDENCE_HALF_LIFE_MS +): number { + const provenance = record.provenance ?? defaultLegacyProvenance(record) + if (provenance.kind === 'user' || halfLifeMs <= 0) return record.confidence + const createdAtMs = Date.parse(record.createdAt) + if (!Number.isFinite(createdAtMs)) return record.confidence + const ageMs = Math.max(0, nowMs - createdAtMs) + return record.confidence * Math.pow(0.5, ageMs / halfLifeMs) +} + +function defaultProvenance(input: MemoryCreateRequest): MemoryProvenance { + return { + kind: 'user', + ...(input.sourceTurnId ? { turnId: input.sourceTurnId } : {}), + origin: 'memory' + } +} + +function defaultLegacyProvenance(record: Pick): MemoryProvenance { + return { + kind: 'user', + ...(record.sourceTurnId ? { turnId: record.sourceTurnId } : {}), + origin: 'legacy' + } +} + +function defaultConfidence(kind: MemoryProvenance['kind']): number { + switch (kind) { + case 'user': return 1 + case 'file': return 0.8 + case 'tool': return 0.7 + case 'web': return 0.5 + case 'inference': return 0.4 + } } /** diff --git a/kun/src/ports/model-client.ts b/kun/src/ports/model-client.ts index 9e5a0dd9a..e9bd0f810 100644 --- a/kun/src/ports/model-client.ts +++ b/kun/src/ports/model-client.ts @@ -48,6 +48,7 @@ export type ModelRequest = { history: TurnItem[] attachments?: ModelInputAttachment[] attachmentTextFallbacks?: ModelTextAttachmentFallback[] + attachmentDocuments?: ModelDocumentAttachment[] tools: ModelToolSpec[] /** * Optional loop-level requirement. The agent loop uses this to keep @@ -94,6 +95,17 @@ export type ModelTextAttachmentFallback = { wasCompressed?: boolean } +export type ModelDocumentAttachment = { + id: string + name: string + mimeType: string + text: string + byteSize: number + pageCount?: number + truncated?: boolean + localFilePath?: string +} + export type ModelToolSpec = { name: string description: string diff --git a/kun/src/ports/tool-host.ts b/kun/src/ports/tool-host.ts index 95765cc56..808d0f4fa 100644 --- a/kun/src/ports/tool-host.ts +++ b/kun/src/ports/tool-host.ts @@ -2,6 +2,7 @@ import type { ApprovalPolicy, SandboxMode } from '../contracts/policy.js' import type { ApprovalRequest } from '../domain/approval.js' import type { TurnItem } from '../contracts/items.js' import type { ModelCapabilityMetadata } from '../contracts/capabilities.js' +import type { ArtifactStore } from '../artifacts/artifact-store.js' import type { UserInputRequest, UserInputResolution @@ -93,6 +94,8 @@ export type ToolHostContext = { sandboxMode?: SandboxMode /** Kun runtime data root; used to allow sandbox-safe reads of background shell output files. */ runtimeDataDir?: string + /** Store used to offload oversized tool results from model context. */ + artifactStore?: ArtifactStore abortSignal: AbortSignal /** Resolves a pending approval with the user's decision. */ awaitApproval: (approval: ApprovalRequest) => Promise<'allow' | 'deny'> diff --git a/kun/src/prompt/kun-system-prompt.ts b/kun/src/prompt/kun-system-prompt.ts index 06a28d11e..8962da77b 100644 --- a/kun/src/prompt/kun-system-prompt.ts +++ b/kun/src/prompt/kun-system-prompt.ts @@ -27,7 +27,7 @@ export const KUN_SYSTEM_PROMPT = [ 'Tool behavior:', '- Use tools when they are available and relevant. Do not claim a file, command, route, or UI state was checked unless it was actually checked.', '- The default built-in coding tool family is `read`, `bash`, `edit`, `write`, `grep`, `find`, and `ls`. Prefer these over ad hoc prose about what you would inspect or change.', - '- Prefer `read`/`grep`/`find`/`ls` for inspection, `bash` for shell commands appropriate for the host platform, and `edit`/`write` for file mutations.', + '- Prefer the most specific advertised tool for the task. Use `read`/`grep`/`find`/`ls` as general inspection fallbacks, `bash` for shell commands appropriate for the host platform, and `edit`/`write` for file mutations.', '- Approval and request_user_input are explicit GUI gates. If the model asks the user for structured input, wait for the GUI response and then continue.', '- Tool results are part of conversation history. Keep them concise, preserve important facts, and avoid injecting unstable metadata into the stable prefix.', '- If a tool is not advertised in the current turn, do not call it.', @@ -57,3 +57,47 @@ export const KUN_SYSTEM_PROMPT = [ '- If a requirement says a capability must not be missing, audit the old surface and prove parity with code paths and tests.', '- A task is complete only when the current code, tests, build, and relevant runtime behavior prove it.' ].join('\n') + +type ToolPreferenceSpec = { + name: string + description: string + providerKind?: string +} + +const SOURCE_EXPLORATION_PATTERN = + /\b(?:code(?:base|graph)?|source|repository|repo|symbol|definition|reference|implementation|dependency|call[ -]?graph|ast)\b/i + +/** + * Keep availability-dependent guidance after the immutable system prefix. + * Tool schemas remain canonically sorted for prompt-cache stability; this + * instruction carries the semantic preference instead of reordering them. + */ +export function buildToolPreferenceInstruction( + tools: readonly ToolPreferenceSpec[] +): string | null { + const mcpTools = tools.filter((tool) => tool.providerKind === 'mcp') + if (mcpTools.length === 0) return null + + const sourceTools = mcpTools.filter((tool) => + SOURCE_EXPLORATION_PATTERN.test(`${tool.name.replace(/[_-]+/g, ' ')} ${tool.description}`) + ) + if (sourceTools.length > 0) { + return [ + `Specialized source-code MCP tools are available for this turn: ${formatToolNames(sourceTools)}.`, + 'For source navigation and structural inspection, prefer a listed MCP tool whose description matches the task before broad `read`/`grep`/`find`/`ls` scans.', + 'Use the built-in inspection tools for unsupported files, narrow fallback checks, and verification.' + ].join(' ') + } + + if (mcpTools.some((tool) => tool.name === 'mcp_search')) { + return 'MCP tool discovery is available through `mcp_search`. When a task may benefit from a specialized external tool, search the MCP catalog before using a general built-in fallback.' + } + + return `Specialized MCP tools are available for this turn: ${formatToolNames(mcpTools)}. Prefer one when its advertised description directly matches the task; otherwise use the built-in tools.` +} + +function formatToolNames(tools: readonly ToolPreferenceSpec[]): string { + const names = tools.slice(0, 8).map((tool) => `\`${tool.name}\``).join(', ') + const remaining = tools.length - 8 + return remaining > 0 ? `${names}, and ${remaining} more` : names +} diff --git a/kun/src/security/secret-store.test.ts b/kun/src/security/secret-store.test.ts new file mode 100644 index 000000000..f708c5d6b --- /dev/null +++ b/kun/src/security/secret-store.test.ts @@ -0,0 +1,178 @@ +import { describe, expect, it, vi } from 'vitest' +import { mkdtemp, rm, readFile, writeFile } from 'node:fs/promises' +import { tmpdir } from 'node:os' +import { join } from 'node:path' +import { randomBytes } from 'node:crypto' +import { createAesEncryptor, createSecretEncryptor, isEncryptedEnvelope } from './secret-store.js' + +describe('createAesEncryptor', () => { + it('round-trips a secret', () => { + const enc = createAesEncryptor(randomBytes(32)) + const blob = enc.encrypt('bearer-token-123') + expect(isEncryptedEnvelope(blob)).toBe(true) + expect(blob).not.toContain('bearer-token-123') + expect(enc.decrypt(blob)).toBe('bearer-token-123') + }) + + it('passes through legacy plaintext on decrypt', () => { + const enc = createAesEncryptor(randomBytes(32)) + expect(enc.decrypt('plain-legacy')).toBe('plain-legacy') + }) + + it('rejects a wrong-size key', () => { + expect(() => createAesEncryptor(randomBytes(16))).toThrow(/32 bytes/) + }) + + it('fails to decrypt a tampered blob', () => { + const enc = createAesEncryptor(randomBytes(32)) + const blob = enc.encrypt('secret') + const tampered = blob.slice(0, -4) + 'AAAA' + expect(() => enc.decrypt(tampered)).toThrow() + }) +}) + +describe('createSecretEncryptor', () => { + it('uses the OS keychain when available (darwin)', async () => { + const store = new Map() + const run = vi.fn(async (command: string, args: string[], input?: string) => { + if (args[0] === 'find-generic-password') { + const v = store.get('k') + return v ? { code: 0, stdout: v, stderr: '' } : { code: 1, stdout: '', stderr: 'not found' } + } + if (args[0] === 'add-generic-password') { + store.set('k', args[args.indexOf('-w') + 1]) + return { code: 0, stdout: '', stderr: '' } + } + return { code: 1, stdout: '', stderr: '' } + }) + const dir = await mkdtemp(join(tmpdir(), 'kun-secret-')) + try { + const result = await createSecretEncryptor({ keyFilePath: join(dir, 'secret.key'), platform: 'darwin', run }) + expect(result.osKeychain).toBe(true) + const blob = result.encryptor.encrypt('tok') + // A second resolve reads the SAME key from the keychain and decrypts. + const again = await createSecretEncryptor({ keyFilePath: join(dir, 'secret.key'), platform: 'darwin', run }) + expect(again.encryptor.decrypt(blob)).toBe('tok') + } finally { + await rm(dir, { recursive: true, force: true }) + } + }) + + it('falls back to a 0600 key file when the keychain is unavailable', async () => { + const dir = await mkdtemp(join(tmpdir(), 'kun-secret-')) + try { + const keyPath = join(dir, 'secret.key') + const result = await createSecretEncryptor({ keyFilePath: keyPath, platform: 'win32' }) + expect(result.osKeychain).toBe(false) + expect(result.reason).toContain('key file') + const blob = result.encryptor.encrypt('tok') + // Persisted key file means a fresh resolve decrypts the same blob. + const again = await createSecretEncryptor({ keyFilePath: keyPath, platform: 'win32' }) + expect(again.encryptor.decrypt(blob)).toBe('tok') + await expect(readFile(keyPath, 'utf8')).resolves.toBeTruthy() + } finally { + await rm(dir, { recursive: true, force: true }) + } + }) + + it('DPAPI-protects the key file on Windows when PowerShell is available', async () => { + // Simulate DPAPI as a reversible wrap (prefix 'DP') so the test exercises + // the protect/store + read/unprotect round-trip without a real keychain. + const run = vi.fn(async (_cmd: string, args: string[], input?: string) => { + const script = args[args.length - 1] + if (script.includes('::Protect')) { + const raw = Buffer.from((input ?? '').trim(), 'base64') + return { code: 0, stdout: Buffer.concat([Buffer.from('DP'), raw]).toString('base64'), stderr: '' } + } + if (script.includes('::Unprotect')) { + const blob = Buffer.from((input ?? '').trim(), 'base64') + if (blob.subarray(0, 2).toString() !== 'DP') return { code: 1, stdout: '', stderr: 'bad' } + return { code: 0, stdout: blob.subarray(2).toString('base64'), stderr: '' } + } + return { code: 1, stdout: '', stderr: '' } + }) + const dir = await mkdtemp(join(tmpdir(), 'kun-secret-')) + try { + const keyPath = join(dir, 'secret.key') + const result = await createSecretEncryptor({ keyFilePath: keyPath, platform: 'win32', run }) + expect(result.osKeychain).toBe(true) + expect(result.reason).toContain('DPAPI') + // The on-disk key file is a DPAPI envelope, not a raw key. + const onDisk = await readFile(keyPath, 'utf8') + expect(onDisk.startsWith('dpapi:v1:')).toBe(true) + const blob = result.encryptor.encrypt('tok') + // A fresh resolve unwraps the same DPAPI-protected key and decrypts. + const again = await createSecretEncryptor({ keyFilePath: keyPath, platform: 'win32', run }) + expect(again.osKeychain).toBe(true) + expect(again.encryptor.decrypt(blob)).toBe('tok') + } finally { + await rm(dir, { recursive: true, force: true }) + } + }) + + it('migrates an existing raw key to DPAPI without changing the encryption key', async () => { + const run = vi.fn(async (_cmd: string, args: string[], input?: string) => { + const script = args[args.length - 1] + if (script.includes('::Protect')) { + const raw = Buffer.from((input ?? '').trim(), 'base64') + return { code: 0, stdout: Buffer.concat([Buffer.from('DP'), raw]).toString('base64'), stderr: '' } + } + if (script.includes('::Unprotect')) { + const blob = Buffer.from((input ?? '').trim(), 'base64') + return { code: 0, stdout: blob.subarray(2).toString('base64'), stderr: '' } + } + return { code: 1, stdout: '', stderr: '' } + }) + const dir = await mkdtemp(join(tmpdir(), 'kun-secret-')) + try { + const keyPath = join(dir, 'secret.key') + const key = randomBytes(32) + const oldEncryptor = createAesEncryptor(key) + const blob = oldEncryptor.encrypt('existing-token') + await writeFile(keyPath, key.toString('base64')) + const migrated = await createSecretEncryptor({ keyFilePath: keyPath, platform: 'win32', run }) + expect(migrated.encryptor.decrypt(blob)).toBe('existing-token') + await expect(readFile(keyPath, 'utf8')).resolves.toMatch(/^dpapi:v1:/) + } finally { + await rm(dir, { recursive: true, force: true }) + } + }) + + it('migrates an existing raw key to the macOS keychain before generating a key', async () => { + let stored = '' + const run = vi.fn(async (_command: string, args: string[]) => { + if (args[0] === 'add-generic-password') { + stored = args[args.indexOf('-w') + 1] + return { code: 0, stdout: '', stderr: '' } + } + return { code: 1, stdout: '', stderr: 'not found' } + }) + const dir = await mkdtemp(join(tmpdir(), 'kun-secret-')) + try { + const keyPath = join(dir, 'secret.key') + const key = randomBytes(32) + const blob = createAesEncryptor(key).encrypt('existing-token') + await writeFile(keyPath, key.toString('base64')) + const migrated = await createSecretEncryptor({ keyFilePath: keyPath, platform: 'darwin', run }) + expect(Buffer.from(stored, 'base64')).toEqual(key) + expect(migrated.encryptor.decrypt(blob)).toBe('existing-token') + } finally { + await rm(dir, { recursive: true, force: true }) + } + }) + + it('falls back to a plain key file on Windows when DPAPI is unavailable', async () => { + const run = vi.fn(async () => ({ code: -1, stdout: '', stderr: 'powershell missing' })) + const dir = await mkdtemp(join(tmpdir(), 'kun-secret-')) + try { + const keyPath = join(dir, 'secret.key') + const result = await createSecretEncryptor({ keyFilePath: keyPath, platform: 'win32', run }) + expect(result.osKeychain).toBe(false) + expect(result.reason).toContain('key file') + const onDisk = await readFile(keyPath, 'utf8') + expect(onDisk.startsWith('dpapi:v1:')).toBe(false) + } finally { + await rm(dir, { recursive: true, force: true }) + } + }) +}) diff --git a/kun/src/security/secret-store.ts b/kun/src/security/secret-store.ts new file mode 100644 index 000000000..d540351d8 --- /dev/null +++ b/kun/src/security/secret-store.ts @@ -0,0 +1,309 @@ +/** + * Secret encryption at rest (review fix: OAuth tokens must not be plaintext). + * + * Provides a symmetric encryptor (AES-256-GCM) whose key is sourced, in order + * of preference, from the OS keychain (macOS `security`, Windows DPAPI via + * PowerShell, Linux `secret-tool`) and otherwise from a per-user key file with + * 0600 permissions. The key-file fallback is a documented, secure-enough + * degradation: the bearer tokens are still encrypted at rest (not stored in + * cleartext JSON), and the key file is owner-only. The command runner and + * platform are injectable so the whole thing is unit-testable without a real + * keychain. + */ + +import { createCipheriv, createDecipheriv, randomBytes } from 'node:crypto' +import { spawn } from 'node:child_process' +import { chmod, mkdir, readFile, rename, rm, writeFile } from 'node:fs/promises' +import { dirname } from 'node:path' + +export type SecretEncryptor = { + encrypt: (plaintext: string) => string + decrypt: (blob: string) => string +} + +const ALGORITHM = 'aes-256-gcm' +const ENVELOPE_PREFIX = 'enc:v1:' + +/** Build an AES-256-GCM encryptor from a 32-byte key. */ +export function createAesEncryptor(key: Buffer): SecretEncryptor { + if (key.length !== 32) throw new Error('encryption key must be 32 bytes') + return { + encrypt: (plaintext: string): string => { + const iv = randomBytes(12) + const cipher = createCipheriv(ALGORITHM, key, iv) + const enc = Buffer.concat([cipher.update(plaintext, 'utf8'), cipher.final()]) + const tag = cipher.getAuthTag() + return `${ENVELOPE_PREFIX}${iv.toString('base64')}:${tag.toString('base64')}:${enc.toString('base64')}` + }, + decrypt: (blob: string): string => { + if (!blob.startsWith(ENVELOPE_PREFIX)) return blob // legacy plaintext + const [, , ivB64, tagB64, dataB64] = blob.split(':') + const iv = Buffer.from(ivB64, 'base64') + const tag = Buffer.from(tagB64, 'base64') + const data = Buffer.from(dataB64, 'base64') + const decipher = createDecipheriv(ALGORITHM, key, iv) + decipher.setAuthTag(tag) + return Buffer.concat([decipher.update(data), decipher.final()]).toString('utf8') + } + } +} + +export function isEncryptedEnvelope(value: string): boolean { + return value.startsWith(ENVELOPE_PREFIX) +} + +export type CommandResult = { code: number; stdout: string; stderr: string } +export type CommandRunner = (command: string, args: string[], input?: string) => Promise + +/** + * Default command runner (spawn, shell:false) used in production so the OS + * keychain / Windows DPAPI path actually runs. Supports feeding `input` over + * stdin so secret material never appears in argv (visible in process listings). + */ +export const defaultSecretCommandRunner: CommandRunner = (command, args, input) => + new Promise((resolve) => { + const maxOutputBytes = 64 * 1024 + let settled = false + const finish = (result: CommandResult): void => { + if (settled) return + settled = true + clearTimeout(timer) + resolve(result) + } + let child: ReturnType + try { + child = spawn(command, args, { shell: false }) + } catch { + resolve({ code: -1, stdout: '', stderr: 'spawn failed' }) + return + } + let stdout = '' + let stderr = '' + const append = (current: string, chunk: Buffer): string => + Buffer.concat([Buffer.from(current, 'utf8'), chunk]).subarray(0, maxOutputBytes).toString('utf8') + child.stdout?.on('data', (d: Buffer) => { stdout = append(stdout, d) }) + child.stderr?.on('data', (d: Buffer) => { stderr = append(stderr, d) }) + const timer = setTimeout(() => { + child.kill('SIGKILL') + finish({ code: -1, stdout, stderr: 'credential helper timed out' }) + }, 10_000) + timer.unref() + child.on('error', () => finish({ code: -1, stdout: '', stderr: 'spawn error' })) + child.on('close', (code) => finish({ code: code ?? -1, stdout, stderr })) + if (input !== undefined) { + child.stdin?.end(input) + } else { + child.stdin?.end() + } + }) + +export type KeyProviderResult = { + encryptor: SecretEncryptor + /** True when the key is protected by the OS keychain (vs the 0600 key file). */ + osKeychain: boolean + /** Human-readable note about how the key is protected (and any degradation). */ + reason: string +} + +const KEYCHAIN_SERVICE = 'kun-secret-key' +const KEYCHAIN_ACCOUNT = 'kun' + +async function tryReadOsKey(platform: NodeJS.Platform, run: CommandRunner): Promise { + try { + if (platform === 'darwin') { + const res = await run('security', ['find-generic-password', '-s', KEYCHAIN_SERVICE, '-a', KEYCHAIN_ACCOUNT, '-w']) + if (res.code === 0 && res.stdout.trim()) return Buffer.from(res.stdout.trim(), 'base64') + } else if (platform === 'linux') { + const res = await run('secret-tool', ['lookup', 'service', KEYCHAIN_SERVICE, 'account', KEYCHAIN_ACCOUNT]) + if (res.code === 0 && res.stdout.trim()) return Buffer.from(res.stdout.trim(), 'base64') + } + } catch { + // tool missing → caller falls back to the key file + } + return null +} + +async function tryWriteOsKey(platform: NodeJS.Platform, run: CommandRunner, key: Buffer): Promise { + const b64 = key.toString('base64') + try { + if (platform === 'darwin') { + const res = await run('security', ['add-generic-password', '-U', '-s', KEYCHAIN_SERVICE, '-a', KEYCHAIN_ACCOUNT, '-w', b64]) + return res.code === 0 + } + if (platform === 'linux') { + const res = await run('secret-tool', ['store', '--label=kun secret key', 'service', KEYCHAIN_SERVICE, 'account', KEYCHAIN_ACCOUNT], b64) + return res.code === 0 + } + } catch { + return false + } + return false +} + +/** Marker for a DPAPI-wrapped key file (Windows). */ +const DPAPI_PREFIX = 'dpapi:v1:' + +/** + * Windows DPAPI key wrapping. The 32-byte AES key is encrypted with the current + * user's DPAPI master key (CurrentUser scope) and the protected blob is what we + * persist — so even the key file cannot be decrypted by another user or off the + * machine. Secret material is fed over stdin, never argv. PowerShell is invoked + * via the injectable runner so this stays unit-testable. + */ +const DPAPI_PROTECT_SCRIPT = + "Add-Type -AssemblyName System.Security; $in=[Console]::In.ReadToEnd().Trim(); " + + "$b=[Convert]::FromBase64String($in); " + + "$p=[System.Security.Cryptography.ProtectedData]::Protect($b,$null,[System.Security.Cryptography.DataProtectionScope]::CurrentUser); " + + "[Console]::Out.Write([Convert]::ToBase64String($p))" +const DPAPI_UNPROTECT_SCRIPT = + "Add-Type -AssemblyName System.Security; $in=[Console]::In.ReadToEnd().Trim(); " + + "$b=[Convert]::FromBase64String($in); " + + "$p=[System.Security.Cryptography.ProtectedData]::Unprotect($b,$null,[System.Security.Cryptography.DataProtectionScope]::CurrentUser); " + + "[Console]::Out.Write([Convert]::ToBase64String($p))" + +async function dpapiProtect(run: CommandRunner, key: Buffer): Promise { + try { + const res = await run('powershell', ['-NoProfile', '-NonInteractive', '-Command', DPAPI_PROTECT_SCRIPT], key.toString('base64')) + if (res.code === 0 && res.stdout.trim()) return res.stdout.trim() + } catch { + // PowerShell missing / blocked → caller falls back to a plain key file. + } + return null +} + +async function dpapiUnprotect(run: CommandRunner, protectedBlobB64: string): Promise { + try { + const res = await run('powershell', ['-NoProfile', '-NonInteractive', '-Command', DPAPI_UNPROTECT_SCRIPT], protectedBlobB64) + if (res.code === 0 && res.stdout.trim()) { + const key = Buffer.from(res.stdout.trim(), 'base64') + return key.length === 32 ? key : null + } + } catch { + // unwrap failed → caller treats the key as unavailable + } + return null +} + +/** Read + DPAPI-unwrap a Windows key file; null when absent/unreadable/unwrappable. */ +async function readDpapiKeyFile(path: string, run: CommandRunner): Promise { + let content: string + try { + content = (await readFile(path, 'utf8')).trim() + } catch { + return null + } + if (!content.startsWith(DPAPI_PREFIX)) return null + return dpapiUnprotect(run, content.slice(DPAPI_PREFIX.length)) +} + +async function readKeyFile(path: string): Promise { + try { + const b64 = (await readFile(path, 'utf8')).trim() + const key = Buffer.from(b64, 'base64') + return key.length === 32 ? key : null + } catch { + return null + } +} + +async function writeKeyFile(path: string, key: Buffer): Promise { + await writeKeyFileContent(path, key.toString('base64')) +} + +/** Write arbitrary key-file content (raw key base64, or a DPAPI envelope) with 0600. */ +async function writeKeyFileContent(path: string, content: string): Promise { + await mkdir(dirname(path), { recursive: true, mode: 0o700 }) + const temporaryPath = `${path}.${process.pid}.${Date.now()}.tmp` + try { + await writeFile(temporaryPath, content, { mode: 0o600 }) + await chmod(temporaryPath, 0o600).catch(() => undefined) + await rename(temporaryPath, path) + await chmod(path, 0o600).catch(() => undefined) + } finally { + await rm(temporaryPath, { force: true }).catch(() => undefined) + } +} + +export type CreateKeyProviderOptions = { + /** Path to the fallback key file (used when the OS keychain is unavailable). */ + keyFilePath: string + platform?: NodeJS.Platform + run?: CommandRunner + /** Disable OS keychain usage (force the key-file fallback). */ + disableOsKeychain?: boolean +} + +/** + * Resolve a secret encryptor. Tries the OS keychain first (storing a random key + * there), then falls back to a 0600 key file. Tokens are always encrypted at + * rest either way. + */ +export async function createSecretEncryptor(options: CreateKeyProviderOptions): Promise { + const platform = options.platform ?? process.platform + const run = options.run + const useKeychain = !options.disableOsKeychain && run && (platform === 'darwin' || platform === 'linux') + // Migration rule: an existing raw key is authoritative because it may be the + // only key capable of decrypting already-persisted OAuth credentials. Never + // generate a fresh OS key before checking it. + const legacyFileKey = await readKeyFile(options.keyFilePath) + + if (useKeychain && run) { + if (legacyFileKey) { + if (await tryWriteOsKey(platform, run, legacyFileKey)) { + await rm(options.keyFilePath, { force: true }).catch(() => undefined) + return { encryptor: createAesEncryptor(legacyFileKey), osKeychain: true, reason: 'existing key migrated to OS keychain' } + } + return { encryptor: createAesEncryptor(legacyFileKey), osKeychain: false, reason: 'OS keychain migration failed; existing 0600 key file preserved' } + } + const existing = await tryReadOsKey(platform, run) + if (existing && existing.length === 32) { + return { encryptor: createAesEncryptor(existing), osKeychain: true, reason: 'key loaded from OS keychain' } + } + const fresh = randomBytes(32) + if (await tryWriteOsKey(platform, run, fresh)) { + return { encryptor: createAesEncryptor(fresh), osKeychain: true, reason: 'new key stored in OS keychain' } + } + } + + // Windows: protect the AES key with DPAPI (CurrentUser) and store the wrapped + // blob in the key file, so the on-disk key cannot be decrypted by another user + // or off the machine. Falls through to a plain key file if PowerShell/DPAPI is + // unavailable (tokens are still AES-encrypted at rest either way). + if (!options.disableOsKeychain && run && platform === 'win32') { + const keyFileText = await readFile(options.keyFilePath, 'utf8').catch(() => '') + const existing = await readDpapiKeyFile(options.keyFilePath, run) + if (existing) { + return { encryptor: createAesEncryptor(existing), osKeychain: true, reason: 'key DPAPI-protected (CurrentUser) in key file' } + } + if (keyFileText.trim().startsWith(DPAPI_PREFIX)) { + throw new Error('existing DPAPI-protected OAuth key could not be decrypted; refusing to replace it') + } + const key = legacyFileKey ?? randomBytes(32) + const wrapped = await dpapiProtect(run, key) + if (wrapped) { + await writeKeyFileContent(options.keyFilePath, `${DPAPI_PREFIX}${wrapped}`) + return { + encryptor: createAesEncryptor(key), + osKeychain: true, + reason: legacyFileKey ? 'existing key migrated to DPAPI protection (CurrentUser)' : 'new key DPAPI-protected (CurrentUser)' + } + } + } + + // Fallback: per-user 0600 key file. + const existingFileKey = legacyFileKey + if (existingFileKey) { + return { + encryptor: createAesEncryptor(existingFileKey), + osKeychain: false, + reason: 'OS keychain unavailable; key loaded from 0600 key file (tokens still encrypted at rest)' + } + } + const fileKey = randomBytes(32) + await writeKeyFile(options.keyFilePath, fileKey) + return { + encryptor: createAesEncryptor(fileKey), + osKeychain: false, + reason: 'OS keychain unavailable; created a new 0600 key file (tokens still encrypted at rest)' + } +} diff --git a/kun/src/security/untrusted-content.ts b/kun/src/security/untrusted-content.ts new file mode 100644 index 000000000..4130d81ad --- /dev/null +++ b/kun/src/security/untrusted-content.ts @@ -0,0 +1,15 @@ +export function wrapUntrustedContent(input: { + content: string + source?: { kind?: string; label?: string } +}): string { + const source = [input.source?.kind, input.source?.label].filter(Boolean).join(':') || 'external' + return [ + ``, + input.content, + '' + ].join('\n') +} + +function escapeAttribute(value: string): string { + return value.replace(/&/g, '&').replace(/"/g, '"').replace(/ { +const DEFAULT_MAX_JSON_BODY_BYTES = 32 * 1024 * 1024 + +export async function readJsonBody(request: Request, maxBytes = DEFAULT_MAX_JSON_BODY_BYTES): Promise { if (request.body === null) return { ok: true, value: {} } - const text = await request.text() + const declaredLength = Number(request.headers.get('content-length')) + if (Number.isFinite(declaredLength) && declaredLength > maxBytes) return bodyTooLarge(maxBytes) + + const reader = request.body.getReader() + const chunks: Uint8Array[] = [] + let totalBytes = 0 + for (;;) { + const { done, value } = await reader.read() + if (done) break + totalBytes += value.byteLength + if (totalBytes > maxBytes) { + await reader.cancel().catch(() => undefined) + return bodyTooLarge(maxBytes) + } + chunks.push(value) + } + const text = Buffer.concat(chunks.map((chunk) => Buffer.from(chunk))).toString('utf8') if (!text) return { ok: true, value: {} } try { return { ok: true, value: JSON.parse(text) } @@ -20,3 +38,10 @@ export async function readJsonBody(request: Request): Promise { + if (!authorize(request, runtime)) return ERRORS.unauthorized() + return mcpOAuthDiagnostics(runtime) + }) + router.add('DELETE', '/v1/mcp/oauth', async (request) => { + if (!authorize(request, runtime)) return ERRORS.unauthorized() + return clearMcpOAuth(runtime) + }) + router.add('DELETE', '/v1/mcp/oauth/:id', async (request, ctx) => { + if (!authorize(request, runtime)) return ERRORS.unauthorized() + return clearMcpOAuth(runtime, ctx.params.id) + }) + router.add('POST', '/v1/mcp/oauth/:id', async (request, ctx) => { + if (!authorize(request, runtime)) return ERRORS.unauthorized() + return authorizeMcpOAuth(runtime, ctx.params.id) + }) router.add('GET', '/v1/skills', async (request) => { if (!authorize(request, runtime)) return ERRORS.unauthorized() return listSkills(runtime) diff --git a/kun/src/server/routes/mcp-oauth.ts b/kun/src/server/routes/mcp-oauth.ts new file mode 100644 index 000000000..17809b62b --- /dev/null +++ b/kun/src/server/routes/mcp-oauth.ts @@ -0,0 +1,18 @@ +import { jsonResponse, type JsonResponse } from '../response.js' +import { ERRORS } from './runtime-error.js' +import type { ServerRuntime } from './server-runtime.js' + +export async function mcpOAuthDiagnostics(runtime: ServerRuntime): Promise { + if (!runtime.mcpOAuth) return ERRORS.unavailable('MCP OAuth diagnostics are not available') + return jsonResponse({ servers: await runtime.mcpOAuth() }) +} + +export async function clearMcpOAuth(runtime: ServerRuntime, serverId?: string): Promise { + if (!runtime.clearMcpOAuth) return ERRORS.unavailable('MCP OAuth credential reset is not available') + return jsonResponse(await runtime.clearMcpOAuth(serverId)) +} + +export async function authorizeMcpOAuth(runtime: ServerRuntime, serverId: string): Promise { + if (!runtime.authorizeMcpOAuth) return ERRORS.unavailable('MCP OAuth authorization is not available') + return jsonResponse(await runtime.authorizeMcpOAuth(serverId)) +} diff --git a/kun/src/server/routes/runtime-info.ts b/kun/src/server/routes/runtime-info.ts index 27bddf52f..f184258a3 100644 --- a/kun/src/server/routes/runtime-info.ts +++ b/kun/src/server/routes/runtime-info.ts @@ -11,6 +11,7 @@ export async function runtimeToolDiagnosticsJsonResponse(runtime: ServerRuntime) return jsonResponse(redactSecrets(await (runtime.toolDiagnostics?.() ?? { providers: [], mcpServers: [], + mcpOAuth: [], webProviders: [], skills: { enabled: false, diff --git a/kun/src/server/routes/server-runtime.ts b/kun/src/server/routes/server-runtime.ts index b7f905f96..22a8643f0 100644 --- a/kun/src/server/routes/server-runtime.ts +++ b/kun/src/server/routes/server-runtime.ts @@ -11,7 +11,12 @@ import type { ToolHost, ToolProviderPolicy } from '../../ports/tool-host.js' import type { RuntimeEventRecorder } from '../../services/runtime-event-recorder.js' import type { LlmDebugRecorder } from '../../services/llm-debug-recorder.js' import type { RuntimeInfoResponse } from '../../contracts/runtime-info.js' -import type { McpServerDiagnostic } from '../../adapters/tool/mcp-tool-provider.js' +import type { + McpOAuthAuthorizeResult, + McpOAuthClearResult, + McpOAuthDiagnostic, + McpServerDiagnostic +} from '../../adapters/tool/mcp-tool-provider.js' import type { McpSearchRuntimeDiagnostic } from '../../adapters/tool/mcp-tool-search.js' import type { WebProviderDiagnostic } from '../../adapters/tool/web-tool-provider.js' import type { ImageGenDiagnostic } from '../../adapters/tool/image-gen-tool-provider.js' @@ -35,6 +40,7 @@ import type { ImmutablePrefix } from '../../cache/immutable-prefix.js' export type RuntimeToolDiagnostics = { providers: ToolProviderPolicy[] mcpServers: McpServerDiagnostic[] + mcpOAuth?: McpOAuthDiagnostic[] mcpSearch?: McpSearchRuntimeDiagnostic webProviders: WebProviderDiagnostic[] skills: SkillRuntimeDiagnostics @@ -111,6 +117,9 @@ export type ServerRuntime = { nowIso: () => string info(): RuntimeInfoResponse toolDiagnostics?(): RuntimeToolDiagnostics | Promise + mcpOAuth?(): McpOAuthDiagnostic[] | Promise + clearMcpOAuth?(serverId?: string): Promise + authorizeMcpOAuth?(serverId: string): Promise skills?(): SkillRuntimeDiagnostics | Promise shutdown?(): Promise } diff --git a/kun/src/server/runtime-factory.ts b/kun/src/server/runtime-factory.ts index 761116428..984ed8c02 100644 --- a/kun/src/server/runtime-factory.ts +++ b/kun/src/server/runtime-factory.ts @@ -16,6 +16,9 @@ import { createAgentSdkRuntime } from '../runtime/agent-sdk/agent-sdk-runtime-fa import { buildGoalLocalTools } from '../adapters/tool/goal-tools.js' import { buildTodoLocalTools } from '../adapters/tool/todo-tools.js' import { LocalToolHost, buildDefaultLocalTools } from '../adapters/tool/local-tool-host.js' +import { createReadArtifactTool } from '../adapters/tool/artifact-tool.js' +import { FileArtifactStore } from '../artifacts/artifact-store.js' +import { createTaskGraphTool } from '../adapters/tool/task-graph-tool.js' import { buildMcpToolProviders } from '../adapters/tool/mcp-tool-provider.js' import { buildMemoryToolProviders } from '../adapters/tool/memory-tool-provider.js' import { buildSkillToolProviders } from '../adapters/tool/skill-tool-provider.js' @@ -81,6 +84,7 @@ import { createChildAgentExecutor } from '../delegation/child-agent-executor.js' import { BackgroundShellRuntime } from '../services/background-shell-runtime.js' import { stopBashSessionById, createBashLocalTool } from '../adapters/tool/builtin-bash-tool.js' import { createBackgroundShellTool } from '../adapters/tool/background-shell-tool.js' +import { createSecretEncryptor, defaultSecretCommandRunner } from '../security/secret-store.js' import type { LocalTool } from '../adapters/tool/local-tool-host.js' export type KunServeRuntimeOptions = { @@ -167,6 +171,7 @@ export async function createKunServeRuntime( ] }) const threadService = new ThreadService({ threadStore, sessionStore, events, ids, nowIso }) + const artifactStore = new FileArtifactStore(join(options.dataDir, 'artifacts'), nowIso) const modelProfiles = modelContextProfilesFromConfig({ contextCompaction: options.contextCompaction, models: options.models @@ -220,9 +225,21 @@ export async function createKunServeRuntime( default: defaultModelClient, providers: providerClients }) + const hasMcpOAuth = Object.values(options.capabilities?.mcp?.servers ?? {}).some((server) => + server.oauth?.enabled !== false && Boolean(server.oauth) && server.transport !== 'stdio' + ) + const oauthEncryptor = hasMcpOAuth + ? (await createSecretEncryptor({ + keyFilePath: join(options.dataDir, 'secret.key'), + run: defaultSecretCommandRunner + })).encryptor + : undefined // Independent I/O; all must still finish before the server listens. const [mcpProviders, skillRuntime] = await Promise.all([ - buildMcpToolProviders(options.capabilities?.mcp), + buildMcpToolProviders(options.capabilities?.mcp, { + oauthStorageDir: join(options.dataDir, 'mcp-oauth'), + ...(oauthEncryptor ? { oauthEncryptor } : {}) + }), SkillRuntime.create(options.capabilities?.skills), seedUsageCarryover({ threadStore, sessionStore, usageService }) ]) @@ -301,6 +318,7 @@ export async function createKunServeRuntime( const musicGenProviders = buildMusicGenToolProviders(options.capabilities?.musicGen, { nowIso }) const videoGenProviders = buildVideoGenToolProviders(options.capabilities?.videoGen, { nowIso }) const computerUseProviders = await buildComputerUseToolProviders(options.capabilities?.computerUse) + const taskGraphTool = createTaskGraphTool({ rootDir: join(options.dataDir, 'task-graphs') }) const baseToolProviders = [ { id: 'builtin', @@ -309,6 +327,13 @@ export async function createKunServeRuntime( available: true, tools: withBackgroundShellTools(buildDefaultLocalTools()) }, + { + id: 'artifacts', + kind: 'built-in' as const, + enabled: true, + available: true, + tools: [createReadArtifactTool()] + }, ...mcpProviders.providers, ...webProviders.providers, ...buildMemoryToolProviders(memoryStore), @@ -359,6 +384,7 @@ export async function createKunServeRuntime( events, ...(options.runtime ? { runtime: options.runtime } : {}), ...(memoryStore ? { memoryStore } : {}), + artifactStore, nowIso }), recordExternalUsage: (threadId, usage) => { @@ -440,6 +466,13 @@ export async function createKunServeRuntime( available: true, tools: buildTodoLocalTools(threadService) }, + { + id: 'planning', + kind: 'built-in' as const, + enabled: true, + available: true, + tools: [taskGraphTool] + }, ...buildDelegationToolProviders(delegationRuntime) ]) const toolHost = new LocalToolHost({ @@ -521,6 +554,7 @@ export async function createKunServeRuntime( ...(options.runtime?.toolArgumentRepair ? { toolArgumentRepair: options.runtime.toolArgumentRepair } : {}), ...(resolvedHooks.length ? { hooks: resolvedHooks } : {}), ...(attachmentStore ? { attachmentStore } : {}), + artifactStore, ...(memoryStore ? { memoryStore } : {}), runtimeDataDir: options.dataDir, onPlanWritten: async ({ threadId, planId, relativePath, markdown }) => { @@ -599,6 +633,7 @@ export async function createKunServeRuntime( toolDiagnostics: async () => ({ providers: registry.diagnostics(), mcpServers: mcpProviders.diagnostics, + mcpOAuth: mcpProviders.oauth, mcpSearch: mcpProviders.search, webProviders: webProviders.diagnostics, skills: skillRuntime.diagnostics(), @@ -613,6 +648,9 @@ export async function createKunServeRuntime( musicGen: musicGenProviders.diagnostics, videoGen: videoGenProviders.diagnostics }), + mcpOAuth: async () => mcpProviders.oauth, + clearMcpOAuth: async (serverId) => mcpProviders.clearOAuthCredentials(serverId), + authorizeMcpOAuth: async (serverId) => mcpProviders.authorizeOAuth(serverId), skills: () => skillRuntime.diagnostics(), shutdown: async () => { try { diff --git a/kun/src/services/usage-service.test.ts b/kun/src/services/usage-service.test.ts index 7bc9aa7cc..27c2295c9 100644 --- a/kun/src/services/usage-service.test.ts +++ b/kun/src/services/usage-service.test.ts @@ -30,6 +30,72 @@ describe('usage cache diagnostics', () => { expect(current.cacheMissReasons).toContain('cold_request') }) + it('explains a hit-rate regression once a thread has a warm baseline', () => { + const usage = new UsageService() + const warm = (hit: number, miss: number) => ({ + promptTokens: hit + miss, + completionTokens: 10, + totalTokens: hit + miss + 10, + cacheHitTokens: hit, + cacheMissTokens: miss, + cacheHitRate: hit / (hit + miss), + turns: 1 + }) + + // Two warm turns at ~90% establish the baseline (no regression yet). + usage.record('thread-r', warm(900, 100), signature) + usage.record('thread-r', warm(900, 100), signature) + // A prefix change collapses the hit rate — should be explained. + const dropped = usage.record('thread-r', warm(50, 950), { + ...signature, + prefixFingerprint: 'prefix-b' + }) + + expect(dropped.cacheMissReasons).toContain('stable_prefix_changed') + expect(dropped.cacheSuggestions?.some((s) => /Cache hit rate dropped/.test(s))).toBe(true) + expect(dropped.cacheSuggestions?.some((s) => /stable system prefix changed/.test(s))).toBe(true) + }) + + it('does not re-announce the same regression every turn (cooldown)', () => { + const usage = new UsageService() + const warm = (hit: number, miss: number) => ({ + promptTokens: hit + miss, + completionTokens: 10, + totalTokens: hit + miss + 10, + cacheHitTokens: hit, + cacheMissTokens: miss, + cacheHitRate: hit / (hit + miss), + turns: 1 + }) + usage.record('thread-c', warm(900, 100), signature) + usage.record('thread-c', warm(900, 100), signature) + const first = usage.record('thread-c', warm(50, 950), { ...signature, prefixFingerprint: 'prefix-b' }) + const second = usage.record('thread-c', warm(50, 950), { ...signature, prefixFingerprint: 'prefix-b' }) + + expect(first.cacheSuggestions?.some((s) => /Cache hit rate dropped/.test(s))).toBe(true) + // The very next turn at the same low rate must NOT repeat the announcement. + expect(second.cacheSuggestions?.some((s) => /Cache hit rate dropped/.test(s))).toBe(false) + }) + + it('starts a fresh baseline when the model changes (no cross-model false regression)', () => { + const usage = new UsageService() + const warm = (hit: number, miss: number) => ({ + promptTokens: hit + miss, + completionTokens: 10, + totalTokens: hit + miss + 10, + cacheHitTokens: hit, + cacheMissTokens: miss, + cacheHitRate: hit / (hit + miss), + turns: 1 + }) + usage.record('thread-m', warm(900, 100), signature) + usage.record('thread-m', warm(900, 100), signature) + // Switch model: the first turn on model-b is cold and has a low hit rate, + // but must not be reported as a regression against model-a's baseline. + const switched = usage.record('thread-m', warm(50, 950), { ...signature, model: 'model-b' }) + expect(switched.cacheSuggestions?.some((s) => /Cache hit rate dropped/.test(s))).toBe(false) + }) + it('surfaces the latest-turn cache diagnostic fields in thread usage', () => { const records: ThreadUsageRecord[] = [ { diff --git a/kun/src/services/usage-service.ts b/kun/src/services/usage-service.ts index 3b92b3b99..ef63f94c1 100644 --- a/kun/src/services/usage-service.ts +++ b/kun/src/services/usage-service.ts @@ -4,6 +4,7 @@ import { diagnoseCacheUsage, type CacheRequestSignature } from '../cache/cache-diagnostics.js' +import { analyzeCacheRegression, cacheRegressionSeverityRank } from '../cache/cache-regression.js' import type { DailyUsageBucket, DailyUsageCounters, @@ -24,6 +25,18 @@ export class UsageService { private readonly counter = new UsageCounter() private readonly cache = new CacheTelemetry() private readonly cacheSignatures = new Map() + /** + * Rolling cacheable-hit-rate history keyed by thread + provider/model/endpoint + * so a model or provider switch starts a FRESH baseline instead of polluting + * the previous one. Cold-start and outliers are absorbed by the analyzer's + * median baseline. + */ + private readonly cacheHitHistory = new Map() + /** + * Cooldown state per history key so one regression isn't re-announced every + * turn: we only re-emit when the cooldown window elapses or severity worsens. + */ + private readonly cacheRegressionCooldown = new Map() record( threadId: string, @@ -50,6 +63,7 @@ export class UsageService { this.cache.reset(threadId) this.cache.ingest(threadId, seeded) this.cacheSignatures.delete(threadId) + this.clearCacheHistory(threadId) return seeded } @@ -70,6 +84,23 @@ export class UsageService { this.cache.reset(threadId) if (threadId === undefined) this.cacheSignatures.clear() else this.cacheSignatures.delete(threadId) + if (threadId === undefined) { + this.cacheHitHistory.clear() + this.cacheRegressionCooldown.clear() + } else { + this.clearCacheHistory(threadId) + } + } + + /** Drop all signature-keyed cache history + cooldown rows for one thread. */ + private clearCacheHistory(threadId: string): void { + const prefix = `${threadId}::` + for (const key of this.cacheHitHistory.keys()) { + if (key === threadId || key.startsWith(prefix)) this.cacheHitHistory.delete(key) + } + for (const key of this.cacheRegressionCooldown.keys()) { + if (key === threadId || key.startsWith(prefix)) this.cacheRegressionCooldown.delete(key) + } } private withCacheDiagnostics( @@ -86,18 +117,66 @@ export class UsageService { ...signature, activeSkillIds: [...signature.activeSkillIds] }) + // Trend-based regression: compare this turn's cacheable hit rate against + // the thread's recent baseline FOR THE SAME provider/model/endpoint so we + // can explain a *drop* (not just the per-turn miss reasons). A cooldown + // stops the same regression from being re-announced every turn. + const historyKey = cacheHistoryKey(threadId, signature) + const history = this.cacheHitHistory.get(historyKey) ?? [] + const regression = analyzeCacheRegression({ + current: diagnostic.cacheableTokenHitRate, + baseline: history, + reasons: diagnostic.reasons, + minBaselineSamples: 2 + }) + const cooldown = this.cacheRegressionCooldown.get(historyKey) ?? { severityRank: 0, turnsSinceEmit: CACHE_REGRESSION_COOLDOWN_TURNS } + cooldown.turnsSinceEmit += 1 + const severityRank = cacheRegressionSeverityRank(regression.severity) + const shouldAnnounce = Boolean(regression.explanation) && ( + cooldown.turnsSinceEmit >= CACHE_REGRESSION_COOLDOWN_TURNS || severityRank > cooldown.severityRank + ) + const suggestions = shouldAnnounce && regression.explanation + ? [regression.explanation, ...diagnostic.suggestions] + : diagnostic.suggestions + if (shouldAnnounce) { + this.cacheRegressionCooldown.set(historyKey, { severityRank, turnsSinceEmit: 0 }) + } else if (regression.severity === 'none') { + // Recovered — clear cooldown so the next genuine drop is announced promptly. + this.cacheRegressionCooldown.set(historyKey, { severityRank: 0, turnsSinceEmit: cooldown.turnsSinceEmit }) + } else { + this.cacheRegressionCooldown.set(historyKey, cooldown) + } + if (typeof diagnostic.cacheableTokenHitRate === 'number') { + this.cacheHitHistory.set(historyKey, [...history, diagnostic.cacheableTokenHitRate].slice(-CACHE_HIT_HISTORY_LIMIT)) + } return { ...usage, cacheableTokenHitRate: diagnostic.cacheableTokenHitRate, totalInputTokenHitRate: diagnostic.totalInputTokenHitRate, cacheMissReasons: diagnostic.reasons, - cacheSuggestions: diagnostic.suggestions + cacheSuggestions: [...new Set(suggestions)] } } } export const MAX_DAILY_USAGE_DAYS = 370 +/** Rolling window of recent cacheable-hit-rate samples kept per thread. */ +const CACHE_HIT_HISTORY_LIMIT = 20 + +/** Turns to wait before re-announcing the same cache regression severity. */ +const CACHE_REGRESSION_COOLDOWN_TURNS = 5 + +/** + * History/cooldown key: per thread AND per provider/model/endpoint so a model + * or provider switch starts a fresh baseline instead of polluting the prior + * one. The prefix fingerprint is intentionally excluded — a prefix change is a + * regression *cause* we want to detect, not a reason to reset the baseline. + */ +function cacheHistoryKey(threadId: string, signature: CacheRequestSignature): string { + return `${threadId}::${signature.providerId}::${signature.model}::${signature.endpointFormat}` +} + const DATE_RE = /^\d{4}-\d{2}-\d{2}$/ export class UsageValidationError extends Error { diff --git a/kun/src/tasks/task-graph.test.ts b/kun/src/tasks/task-graph.test.ts new file mode 100644 index 000000000..96fba4573 --- /dev/null +++ b/kun/src/tasks/task-graph.test.ts @@ -0,0 +1,117 @@ +import { describe, expect, it } from 'vitest' +import { TaskGraph } from './task-graph.js' + +describe('TaskGraph', () => { + it('marks tasks ready only when dependencies have succeeded', () => { + const g = new TaskGraph({ concurrency: 2 }) + g.add({ id: 'a', title: 'A' }) + g.add({ id: 'b', title: 'B', dependsOn: ['a'] }) + g.reconcile() + expect(g.get('a')?.state).toBe('ready') + expect(g.get('b')?.state).toBe('pending') + g.markRunning('a'); g.markSucceeded('a') + expect(g.get('b')?.state).toBe('ready') + }) + + it('respects priority and concurrency in nextRunnable', () => { + const g = new TaskGraph({ concurrency: 2 }) + g.add({ id: 'low', title: 'low', priority: 1 }) + g.add({ id: 'high', title: 'high', priority: 10 }) + g.add({ id: 'mid', title: 'mid', priority: 5 }) + const batch = g.nextRunnable() + expect(batch.map((t) => t.id)).toEqual(['high', 'mid']) + }) + + it('does not exceed the concurrency limit', () => { + const g = new TaskGraph({ concurrency: 1 }) + g.add({ id: 'a', title: 'A' }) + g.add({ id: 'b', title: 'B' }) + const first = g.nextRunnable() + expect(first).toHaveLength(1) + g.markRunning(first[0].id) + expect(g.nextRunnable()).toHaveLength(0) + }) + + it('retries a failed task until attempts are exhausted', () => { + const g = new TaskGraph() + g.add({ id: 'a', title: 'A', maxAttempts: 2 }) + g.reconcile() + g.markRunning('a') + expect(g.markFailed('a', 'boom').retried).toBe(true) + expect(g.get('a')?.state).toBe('ready') + g.markRunning('a') + expect(g.markFailed('a', 'boom again').retried).toBe(false) + expect(g.get('a')?.state).toBe('failed') + }) + + it('blocks dependents when a dependency fails terminally', () => { + const g = new TaskGraph() + g.add({ id: 'a', title: 'A', maxAttempts: 1 }) + g.add({ id: 'b', title: 'B', dependsOn: ['a'] }) + g.reconcile() + g.markRunning('a') + g.markFailed('a', 'x') + expect(g.get('b')?.state).toBe('blocked') + }) + + it('detects dependency cycles', () => { + const g = new TaskGraph() + g.add({ id: 'a', title: 'A', dependsOn: ['c'] }) + g.add({ id: 'b', title: 'B', dependsOn: ['a'] }) + expect(() => g.add({ id: 'c', title: 'C', dependsOn: ['b'] })).toThrow(/dependency cycle/) + expect(g.get('c')).toBeUndefined() + expect(g.detectCycle()).toBeNull() + }) + + it('pause/resume removes and restores a task from scheduling', () => { + const g = new TaskGraph() + g.add({ id: 'a', title: 'A' }) + g.pause('a') + expect(g.nextRunnable()).toHaveLength(0) + g.resume('a') + expect(g.nextRunnable().map((t) => t.id)).toEqual(['a']) + }) + + it('round-trips through JSON', () => { + const g = new TaskGraph({ concurrency: 3 }) + g.add({ id: 'a', title: 'A', tokenBudget: 1000, worktree: '/wt/a' }) + const restored = TaskGraph.fromJSON(g.toJSON()) + expect(restored.get('a')).toMatchObject({ tokenBudget: 1000, worktree: '/wt/a' }) + expect(restored.toJSON().concurrency).toBe(3) + }) + + it('reports completion when all tasks are terminal', () => { + const g = new TaskGraph() + g.add({ id: 'a', title: 'A' }) + g.reconcile(); g.markRunning('a'); g.markSucceeded('a') + expect(g.isComplete()).toBe(true) + }) + + it('does not rewrite a completed task when cancellation is replayed', () => { + const g = new TaskGraph() + g.add({ id: 'a', title: 'A' }) + g.reconcile(); g.markRunning('a'); g.markSucceeded('a') + g.cancel('a') + expect(g.get('a')?.state).toBe('succeeded') + }) + + it('blocks a task with a missing dependency and treats blocked as terminal', () => { + const g = new TaskGraph() + g.add({ id: 'b', title: 'B', dependsOn: ['does-not-exist'] }) + g.reconcile() + expect(g.get('b')?.state).toBe('blocked') + expect(g.isComplete()).toBe(true) + }) + + it('applies exponential backoff between retries via the injected clock', () => { + let now = 1000 + const g = new TaskGraph({ now: () => now, retryBaseDelayMs: 100, retryMaxDelayMs: 10_000 }) + g.add({ id: 'a', title: 'A', maxAttempts: 3 }) + g.reconcile(); g.markRunning('a'); g.markFailed('a', 'boom') + // Immediately after failure the task is ready but gated by backoff. + expect(g.get('a')?.state).toBe('ready') + expect(g.nextRunnable()).toHaveLength(0) + now += 100 + expect(g.nextRunnable().map((t) => t.id)).toEqual(['a']) + }) +}) diff --git a/kun/src/tasks/task-graph.ts b/kun/src/tasks/task-graph.ts new file mode 100644 index 000000000..dab34161e --- /dev/null +++ b/kun/src/tasks/task-graph.ts @@ -0,0 +1,269 @@ +/** + * Persistent task graph (P1 #2). + * + * More than a flat todo list or a single useWorktree flag: tasks form a DAG with + * dependencies, priority, a concurrency limit, pause/resume, failure retry, a + * per-task resource budget, and an optional per-task worktree. This is the pure + * scheduling core — it computes the ready set (respecting deps + concurrency), + * applies state transitions, handles retry with backoff bookkeeping, and detects + * cycles. Persistence and actual execution live in a thin layer on top. + */ + +export type TaskState = 'pending' | 'ready' | 'running' | 'blocked' | 'paused' | 'succeeded' | 'failed' | 'cancelled' + +export type TaskNode = { + id: string + title: string + /** Ids this task depends on; it cannot start until all have succeeded. */ + dependsOn: string[] + /** Higher runs first among ready tasks. Default 0. */ + priority: number + state: TaskState + attempts: number + maxAttempts: number + /** Optional per-task worktree path. */ + worktree?: string + /** Optional resource budget (tokens) for the task. */ + tokenBudget?: number + lastError?: string + /** Earliest time (ms epoch) the task may be retried after a failure. */ + nextAttemptAt?: number +} + +export type TaskGraphData = { + tasks: Record + /** Max tasks allowed in `running` at once. */ + concurrency: number +} + +export type AddTaskInput = { + id: string + title: string + dependsOn?: string[] + priority?: number + maxAttempts?: number + worktree?: string + tokenBudget?: number +} + +export type TaskGraphOptions = { + concurrency?: number + now?: () => number + /** First retry backoff; doubles per attempt up to retryMaxDelayMs. Default 1000. */ + retryBaseDelayMs?: number + retryMaxDelayMs?: number +} + +export class TaskGraph { + private readonly tasks = new Map() + private concurrency: number + private readonly now: () => number + private readonly retryBaseDelayMs: number + private readonly retryMaxDelayMs: number + + constructor(input: TaskGraphOptions = {}) { + this.concurrency = Math.max(1, input.concurrency ?? 1) + this.now = input.now ?? (() => Date.now()) + this.retryBaseDelayMs = input.retryBaseDelayMs ?? 1_000 + this.retryMaxDelayMs = input.retryMaxDelayMs ?? 30_000 + } + + add(input: AddTaskInput): TaskNode { + if (this.tasks.has(input.id)) throw new Error(`duplicate task id: ${input.id}`) + const node: TaskNode = { + id: input.id, + title: input.title, + dependsOn: [...(input.dependsOn ?? [])], + priority: input.priority ?? 0, + state: 'pending', + attempts: 0, + maxAttempts: Math.max(1, input.maxAttempts ?? 1), + ...(input.worktree ? { worktree: input.worktree } : {}), + ...(input.tokenBudget ? { tokenBudget: input.tokenBudget } : {}) + } + this.tasks.set(input.id, node) + const cycle = this.detectCycle() + if (cycle) { + this.tasks.delete(input.id) + throw new Error(`dependency cycle: ${cycle.join(' -> ')}`) + } + return node + } + + setConcurrency(limit: number): void { + this.concurrency = Math.max(1, limit) + } + + get(id: string): TaskNode | undefined { + return this.tasks.get(id) + } + + list(): TaskNode[] { + return [...this.tasks.values()] + } + + /** Detect a dependency cycle; returns the cycle path or null when acyclic. */ + detectCycle(): string[] | null { + const visited = new Set() + const stack = new Set() + const path: string[] = [] + const visit = (id: string): string[] | null => { + if (stack.has(id)) return [...path.slice(path.indexOf(id)), id] + if (visited.has(id)) return null + visited.add(id) + stack.add(id) + path.push(id) + for (const dep of this.tasks.get(id)?.dependsOn ?? []) { + if (!this.tasks.has(dep)) continue + const cycle = visit(dep) + if (cycle) return cycle + } + stack.delete(id) + path.pop() + return null + } + for (const id of this.tasks.keys()) { + const cycle = visit(id) + if (cycle) return cycle + } + return null + } + + private depsSucceeded(node: TaskNode): boolean { + return node.dependsOn.every((dep) => this.tasks.get(dep)?.state === 'succeeded') + } + + private depsUnsatisfiable(node: TaskNode): boolean { + return node.dependsOn.some((dep) => { + const depNode = this.tasks.get(dep) + // A missing dependency can never be satisfied → the task is blocked, not + // left pending forever. + if (!depNode) return true + return depNode.state === 'failed' || depNode.state === 'cancelled' || depNode.state === 'blocked' + }) + } + + /** + * Recompute derived states: a pending task whose deps all succeeded becomes + * `ready`; one with a failed/cancelled/blocked/missing dep becomes `blocked`. + * Running/paused/terminal states are left untouched. + */ + reconcile(): void { + for (const node of this.tasks.values()) { + if (node.state !== 'pending' && node.state !== 'ready' && node.state !== 'blocked') continue + if (this.depsUnsatisfiable(node)) { + node.state = 'blocked' + } else if (this.depsSucceeded(node)) { + node.state = 'ready' + } else { + node.state = 'pending' + } + } + } + + /** Number of tasks currently running. */ + runningCount(): number { + return this.list().filter((node) => node.state === 'running').length + } + + /** + * The next batch of tasks to start: ready tasks (deps satisfied, not paused), + * highest priority first, bounded by remaining concurrency. + */ + nextRunnable(): TaskNode[] { + this.reconcile() + const slots = this.concurrency - this.runningCount() + if (slots <= 0) return [] + return this.list() + .filter((node) => node.state === 'ready') + .filter((node) => node.nextAttemptAt === undefined || node.nextAttemptAt <= this.now()) + .sort((a, b) => (b.priority - a.priority) || a.id.localeCompare(b.id)) + .slice(0, slots) + } + + markRunning(id: string): void { + const node = this.require(id) + if (node.state !== 'ready') throw new Error(`task ${id} is not ready (state: ${node.state})`) + node.state = 'running' + node.attempts += 1 + node.nextAttemptAt = undefined + } + + markSucceeded(id: string): void { + const node = this.require(id) + if (node.state !== 'running') throw new Error(`task ${id} is not running (state: ${node.state})`) + node.state = 'succeeded' + node.lastError = undefined + this.reconcile() + } + + /** + * Fail a running task. Retries (back to `ready`) while attempts remain; + * otherwise terminal `failed`, which cascades dependents to `blocked` on the + * next reconcile. + */ + markFailed(id: string, error: string): { retried: boolean } { + const node = this.require(id) + if (node.state !== 'running') throw new Error(`task ${id} is not running (state: ${node.state})`) + node.lastError = error + if (node.attempts < node.maxAttempts) { + node.state = 'ready' + // Exponential backoff so a flapping task does not hot-loop; nextRunnable + // skips it until nextAttemptAt. + const delay = Math.min(this.retryMaxDelayMs, this.retryBaseDelayMs * 2 ** (node.attempts - 1)) + node.nextAttemptAt = this.now() + delay + return { retried: true } + } + node.state = 'failed' + this.reconcile() + return { retried: false } + } + + pause(id: string): void { + const node = this.require(id) + if (node.state === 'succeeded' || node.state === 'failed' || node.state === 'cancelled') return + node.state = 'paused' + } + + resume(id: string): void { + const node = this.require(id) + if (node.state !== 'paused') return + node.state = 'pending' + this.reconcile() + } + + cancel(id: string): void { + const node = this.require(id) + if (node.state === 'succeeded' || node.state === 'failed' || node.state === 'cancelled') return + node.state = 'cancelled' + this.reconcile() + } + + /** True when every task reached a terminal state (blocked is terminal: its deps can never satisfy). */ + isComplete(): boolean { + this.reconcile() + return this.list().every((node) => ['succeeded', 'failed', 'cancelled', 'blocked'].includes(node.state)) + } + + toJSON(): TaskGraphData { + const tasks: Record = {} + for (const [id, node] of this.tasks) tasks[id] = { ...node, dependsOn: [...node.dependsOn] } + return { tasks, concurrency: this.concurrency } + } + + static fromJSON(data: TaskGraphData): TaskGraph { + const graph = new TaskGraph({ concurrency: data.concurrency }) + for (const node of Object.values(data.tasks)) { + graph.tasks.set(node.id, { ...node, dependsOn: [...node.dependsOn] }) + } + const cycle = graph.detectCycle() + if (cycle) throw new Error(`dependency cycle: ${cycle.join(' -> ')}`) + return graph + } + + private require(id: string): TaskNode { + const node = this.tasks.get(id) + if (!node) throw new Error(`unknown task: ${id}`) + return node + } +} diff --git a/kun/tests/attachment-store.test.ts b/kun/tests/attachment-store.test.ts index 1baadc165..6bb29ff99 100644 --- a/kun/tests/attachment-store.test.ts +++ b/kun/tests/attachment-store.test.ts @@ -79,11 +79,17 @@ describe('Attachment store and multimodal input', () => { }) }) + it('rejects attachment ids that could escape the store directory', async () => { + const store = createStore() + await expect(store.get('../outside')).resolves.toBeNull() + await expect(store.resolveContent('..\\outside', {})).rejects.toThrow(/invalid attachment id/) + }) + it('rejects unsupported MIME, size, and dimensions', async () => { await expect(createStore().create({ - name: 'bad.txt', + name: 'bad.bin', data: Buffer.from('nope'), - mimeType: 'text/plain' + mimeType: 'application/octet-stream' })).rejects.toThrow(/unsupported/) await expect(createStore({ maxImageBytes: 10 }).create({ @@ -172,6 +178,21 @@ describe('Attachment store and multimodal input', () => { expect(await readJson(diagnostics)).toMatchObject({ enabled: true, count: 1 }) }) + it('rejects malformed base64 attachment uploads', async () => { + const h = buildHarness() + h.runtime.attachmentStore = createStore() + const response = await dispatchRequest( + h.router, + new Request('http://localhost/v1/attachments', { + method: 'POST', + headers: { authorization: 'Bearer tok-1', 'content-type': 'application/json' }, + body: JSON.stringify({ name: 'shot.png', dataBase64: 'not-base64!' }) + }) + ) + expect(response.status).toBe(400) + expect(await readJson(response)).toMatchObject({ message: 'attachment data is not valid base64' }) + }) + it('resolves image attachments for vision models and text fallbacks for text-only models', async () => { const store = createStore() const workspace = join(dir, 'workspace') @@ -466,6 +487,154 @@ describe('Attachment store and multimodal input', () => { expect(body?.messages?.[0]?.content).toContain('```base64\nYWJj\n```') }) + it('stores PDF document attachments with extracted text', async () => { + const store = createStore() + const doc = await store.create({ + name: 'spec.pdf', + data: Buffer.from('%PDF-1.7\nbinary-body'), + mimeType: 'application/pdf', + documentText: 'Hello from the PDF body.', + pageCount: 3, + threadId: 'thr_1' + }) + + expect(doc).toMatchObject({ + kind: 'document', + mimeType: 'application/pdf', + documentText: 'Hello from the PDF body.', + pageCount: 3 + }) + await expect(store.resolveContent(doc.id, { threadId: 'thr_1' })).resolves.toMatchObject({ + kind: 'document', + documentText: 'Hello from the PDF body.' + }) + }) + + it('decodes and truncates text-like document attachments', async () => { + const store = createStore({ maxDocumentTextChars: 5 }) + const doc = await store.create({ + name: 'notes.md', + data: Buffer.from('\uFEFF0123456789', 'utf8'), + mimeType: 'text/markdown', + threadId: 'thr_1' + }) + + expect(doc).toMatchObject({ + kind: 'document', + documentText: '01234', + truncated: true + }) + }) + + it('rejects document attachments with disallowed MIME or missing extracted text', async () => { + await expect(createStore().create({ + name: 'archive.zip', + data: Buffer.from('PK\u0003\u0004nope'), + mimeType: 'application/zip' + })).rejects.toThrow(/unsupported attachment type/) + + await expect(createStore().create({ + name: 'scan.pdf', + data: Buffer.from('%PDF-1.4 no extractable text'), + mimeType: 'application/pdf' + })).rejects.toThrow(/document text is required/) + }) + + it('injects document attachments into model requests as text context', async () => { + const store = createStore() + const doc = await store.create({ + name: 'spec.pdf', + data: Buffer.from('%PDF-1.7 body'), + mimeType: 'application/pdf', + documentText: 'The deployment runbook lives here.', + pageCount: 2, + threadId: 'thr_1', + workspace: '/tmp/ws' + }) + const seenRequests: ModelRequest[] = [] + const model: ModelClient = { + provider: 'fake', + model: 'fake', + async *stream(request) { + seenRequests.push(request) + yield { kind: 'completed', stopReason: 'stop' } + } + } + const h = makeHarness(model, { + attachmentStore: store, + modelCapabilities: () => ({ ...visionCapabilities(), inputModalities: ['text'] }) + }) + await bootstrapThread(h, { + workspace: '/tmp/ws', + request: { prompt: 'summarize', attachmentIds: [doc.id], model: 'text-only' } + }) + + expect(await h.loop.runTurn(h.threadId, h.turnId)).toBe('completed') + expect(seenRequests.at(-1)?.attachments).toBeUndefined() + expect(seenRequests.at(-1)?.attachmentDocuments?.[0]).toMatchObject({ + id: doc.id, + mimeType: 'application/pdf', + text: 'The deployment runbook lives here.', + pageCount: 2 + }) + }) + + it('formats document attachments as untrusted text for the model', async () => { + let body: { messages?: Array<{ role: string; content: unknown }> } | undefined + const client = new CompatModelClient({ + baseUrl: 'https://model.example.test', + apiKey: '', + model: 'text-model', + nonStreaming: true, + fetchImpl: async (_url, init) => { + body = JSON.parse(String(init?.body)) + return new Response(JSON.stringify({ + id: 'cmpl_1', + model: 'text-model', + choices: [{ index: 0, finish_reason: 'stop', message: { role: 'assistant', content: 'ok' } }] + }), { status: 200, headers: { 'content-type': 'application/json' } }) + } + }) + + for await (const _chunk of client.stream({ + threadId: 'thr_1', + turnId: 'turn_1', + model: 'text-model', + prefix: [], + history: [{ + id: 'item_user', + threadId: 'thr_1', + turnId: 'turn_1', + role: 'user', + status: 'completed', + createdAt: 'now', + finishedAt: 'now', + kind: 'user_message', + text: 'summarize' + }], + attachmentDocuments: [{ + id: 'att_doc', + name: 'spec.pdf', + mimeType: 'application/pdf', + text: 'Runbook contents.', + byteSize: 42, + pageCount: 2, + localFilePath: '/tmp/picked/spec.pdf' + }], + tools: [], + abortSignal: new AbortController().signal + })) { + // drain stream + } + + const content = String(body?.messages?.[0]?.content ?? '') + expect(content).toContain('[Attached document]') + expect(content).toContain('Name: spec.pdf') + expect(content).toContain('Pages: 2') + expect(content).toContain('') + expect(content).toContain('Runbook contents.') + }) + function createStore(overrides: Partial = {}) { return new FileAttachmentStore({ rootDir: join(dir, 'attachments'), diff --git a/kun/tests/builtin-tools.test.ts b/kun/tests/builtin-tools.test.ts index e8dbaba14..e2f01c6b1 100644 --- a/kun/tests/builtin-tools.test.ts +++ b/kun/tests/builtin-tools.test.ts @@ -341,7 +341,7 @@ describe('Kun built-in tools', () => { it('exposes pi-style coding and read-only tool groups', () => { expect(buildCodingBuiltinLocalTools().map((tool) => tool.name)).toEqual(['read', 'bash', 'edit', 'write']) - expect(buildReadOnlyBuiltinLocalTools().map((tool) => tool.name)).toEqual(['read', 'grep', 'find', 'ls']) + expect(buildReadOnlyBuiltinLocalTools().map((tool) => tool.name)).toEqual(['read', 'grep', 'find', 'ls', 'repo_map']) }) it('supports pi-style configurable built-in tool factory APIs', async () => { @@ -352,7 +352,18 @@ describe('Kun built-in tools', () => { ls: { defaultLimit: 1 }, bash: { defaultTimeoutSeconds: 5 } }) - expect(Object.keys(toolRecord).sort()).toEqual(['bash', 'edit', 'find', 'grep', 'ls', 'lsp', 'read', 'write']) + expect(Object.keys(toolRecord).sort()).toEqual([ + 'bash', + 'edit', + 'find', + 'grep', + 'ls', + 'lsp', + 'read', + 'repo_map', + 'verify_changes', + 'write' + ]) await writeFile(join(workspace, 'limited.txt'), 'one\ntwo\nthree\n', 'utf8') const customHost = new LocalToolHost({ tools: [toolRecord.read, toolRecord.ls] }) @@ -369,13 +380,35 @@ describe('Kun built-in tools', () => { expect(defaultGrepLocalToolOperations).toEqual({}) expect(defaultLsLocalToolOperations.readdir).toBeTypeOf('function') expect(createCodingTools().map((tool) => tool.name)).toEqual(['read', 'bash', 'edit', 'write']) - expect(createReadOnlyTools().map((tool) => tool.name)).toEqual(['read', 'grep', 'find', 'ls']) + expect(createReadOnlyTools().map((tool) => tool.name)).toEqual(['read', 'grep', 'find', 'ls', 'repo_map']) expect(createCodingToolDefinitions().map((tool) => tool.name)).toEqual(['read', 'bash', 'edit', 'write']) - expect(createReadOnlyToolDefinitions().map((tool) => tool.name)).toEqual(['read', 'grep', 'find', 'ls']) + expect(createReadOnlyToolDefinitions().map((tool) => tool.name)).toEqual(['read', 'grep', 'find', 'ls', 'repo_map']) const allTools = createAllTools() const allDefinitions = createAllToolDefinitions() - expect(Object.keys(allTools).sort()).toEqual(['bash', 'edit', 'find', 'grep', 'ls', 'lsp', 'read', 'write']) - expect(Object.keys(allDefinitions).sort()).toEqual(['bash', 'edit', 'find', 'grep', 'ls', 'lsp', 'read', 'write']) + expect(Object.keys(allTools).sort()).toEqual([ + 'bash', + 'edit', + 'find', + 'grep', + 'ls', + 'lsp', + 'read', + 'repo_map', + 'verify_changes', + 'write' + ]) + expect(Object.keys(allDefinitions).sort()).toEqual([ + 'bash', + 'edit', + 'find', + 'grep', + 'ls', + 'lsp', + 'read', + 'repo_map', + 'verify_changes', + 'write' + ]) expect(createReadTool).toBe(createReadLocalTool) expect(createReadToolDefinition).toBe(createReadLocalTool) expect(createWriteTool).toBeTypeOf('function') diff --git a/kun/tests/child-agent-executor.test.ts b/kun/tests/child-agent-executor.test.ts index 3aab53953..b15c37896 100644 --- a/kun/tests/child-agent-executor.test.ts +++ b/kun/tests/child-agent-executor.test.ts @@ -87,6 +87,56 @@ describe('child agent executor', () => { expect(seen[0]?.tools).toEqual([]) }) + it('returns bounded tool evidence when the contract requests it', async () => { + let step = 0 + const evidenceModel: ModelClient = { + provider: 'evidence-test', + model: 'evidence-test', + async *stream(): AsyncIterable { + step += 1 + if (step === 1) { + yield { + kind: 'tool_call_complete', + callId: 'call_inspect', + toolName: 'inspect', + arguments: { path: 'src/index.ts' } + } + yield { kind: 'completed', stopReason: 'tool_calls' } + return + } + yield { kind: 'assistant_text_delta', text: 'Inspection complete.' } + yield { kind: 'completed', stopReason: 'stop' } + } + } + const inspect = LocalToolHost.defineTool({ + name: 'inspect', + description: 'Inspect a source file.', + inputSchema: { type: 'object' }, + policy: 'auto', + execute: async () => ({ output: { ok: true } }) + }) + const executor = createChildAgentExecutor({ + model: evidenceModel, + toolHost: new LocalToolHost({ tools: [inspect] }), + prefix: createImmutablePrefix({ systemPrompt: 'child system' }), + defaultModel: 'evidence-test', + nowIso: () => '2026-06-03T00:00:00.000Z' + }) + + const result = await executor({ + childId: 'child_evidence', + parentThreadId: 'thr_parent', + parentTurnId: 'turn_parent', + prompt: 'Inspect the entry point', + toolPolicy: 'inherit', + returnFormat: 'evidence', + signal: new AbortController().signal + }) + + expect(result.summary).toBe('Inspection complete.') + expect(result.evidence).toEqual(['inspect src/index.ts: completed']) + }) + it('fails the child run when the child loop cannot produce a completed turn', async () => { const executor = createChildAgentExecutor({ model: model([{ kind: 'error', message: 'model failed', code: 'bad_model' }]), @@ -137,7 +187,7 @@ describe('child agent executor', () => { }) const toolNames = (seen[0]?.tools ?? []).map((tool) => tool.name).sort() - expect(toolNames).toEqual(['find', 'grep', 'ls', 'read']) + expect(toolNames).toEqual(['find', 'grep', 'ls', 'read', 'repo_map']) expect(seen[0]?.history?.[0]).toMatchObject({ kind: 'user_message', text: 'Read-only review.\n\nInvestigate the bug' diff --git a/kun/tests/cli-agent.test.ts b/kun/tests/cli-agent.test.ts index 96fc952a4..964875cb3 100644 --- a/kun/tests/cli-agent.test.ts +++ b/kun/tests/cli-agent.test.ts @@ -240,6 +240,31 @@ describe('Kun agent CLI commands', () => { expect(parsed.items.some((item) => item.kind === 'assistant_text')).toBe(true) }) + it('streams a stable JSONL envelope for headless runs', async () => { + const c = capture({ createRuntime: fakeRuntime() }) + const code = await runAgentCommand('run', [ + '--data-dir', + dataDir, + '--prompt', + 'hello', + '--jsonl' + ], c.io) + + expect(code).toBe(ServeExitCode.ok) + const lines = c.stdout.trim().split('\n').map((line) => JSON.parse(line) as Record) + expect(lines[0]).toMatchObject({ type: 'run_started' }) + expect(lines.at(-1)).toMatchObject({ type: 'run_finished', status: 'completed' }) + }) + + it('rejects combining JSON and JSONL output modes', async () => { + const c = capture({ createRuntime: fakeRuntime() }) + const code = await runAgentCommand('run', [ + '--data-dir', dataDir, '--prompt', 'hello', '--json', '--jsonl' + ], c.io) + expect(code).toBe(ServeExitCode.usage) + expect(c.stderr).toContain('mutually exclusive') + }) + it('returns runtime failures from one-shot runs', async () => { const c = capture({ createRuntime: fakeRuntime({ throwRun: true }) }) const code = await runAgentCommand('run', [ diff --git a/kun/tests/create-plan-tool.test.ts b/kun/tests/create-plan-tool.test.ts index 723c47619..566a2e293 100644 --- a/kun/tests/create-plan-tool.test.ts +++ b/kun/tests/create-plan-tool.test.ts @@ -92,6 +92,7 @@ describe('create_plan tool: advertisement', () => { 'grep', 'find', 'ls', + 'repo_map', 'user_input', 'request_user_input', CREATE_PLAN_TOOL_NAME @@ -104,7 +105,7 @@ describe('create_plan tool: advertisement', () => { const tools = await host.listTools(buildContext({ threadMode: 'plan' })) const names = tools.map((tool) => tool.name) - expect(names).toEqual(['read', 'grep', 'find', 'ls', CREATE_PLAN_TOOL_NAME]) + expect(names).toEqual(['read', 'grep', 'find', 'ls', 'repo_map', CREATE_PLAN_TOOL_NAME]) }) it('keeps the normal agent-mode default tool advertisement unchanged', async () => { diff --git a/kun/tests/delegation-runtime.test.ts b/kun/tests/delegation-runtime.test.ts index afa0ffec8..9f6608f3b 100644 --- a/kun/tests/delegation-runtime.test.ts +++ b/kun/tests/delegation-runtime.test.ts @@ -81,6 +81,80 @@ describe('DelegationRuntime', () => { })).rejects.toThrow(/budget/) }) + it('enforces per-child token and time budgets', async () => { + const externalUsage: unknown[] = [] + const tokenRuntime = createRuntime({ + recordExternalUsage: (_threadId, usage) => externalUsage.push(usage), + executor: async () => ({ + summary: 'used too many tokens', + usage: { promptTokens: 8, completionTokens: 4, totalTokens: 12 } + }) + }) + const tokenLimited = await tokenRuntime.runChild({ + parentThreadId: 'thr_tokens', + parentTurnId: 'turn_tokens', + prompt: 'bounded task', + tokenBudget: 10, + signal: new AbortController().signal + }) + expect(tokenLimited).toMatchObject({ + status: 'failed', + tokenBudget: 10, + budgetExceeded: 'token', + usage: { totalTokens: 12 } + }) + expect(tokenLimited.error).toContain('12 > 10') + expect(externalUsage).toHaveLength(1) + + const timeRuntime = createRuntime({ + executor: async ({ signal }) => new Promise((_resolve, reject) => { + signal.addEventListener('abort', () => reject(signal.reason), { once: true }) + }) + }) + const timeLimited = await timeRuntime.runChild({ + parentThreadId: 'thr_time', + parentTurnId: 'turn_time', + prompt: 'slow task', + timeBudgetMs: 10, + signal: new AbortController().signal + }) + expect(timeLimited).toMatchObject({ + status: 'failed', + timeBudgetMs: 10, + budgetExceeded: 'time' + }) + expect(timeLimited.error).toContain('10ms') + }) + + it('validates evidence-return contracts', async () => { + const withEvidence = createRuntime({ + executor: async () => ({ summary: 'done', evidence: ['read src/index.ts', 'ran unit tests'] }) + }) + await expect(withEvidence.runChild({ + parentThreadId: 'thr_evidence', + parentTurnId: 'turn_evidence', + prompt: 'investigate', + returnFormat: 'evidence', + signal: new AbortController().signal + })).resolves.toMatchObject({ + status: 'completed', + returnFormat: 'evidence', + evidence: ['read src/index.ts', 'ran unit tests'] + }) + + const withoutEvidence = createRuntime({ executor: async () => ({ summary: 'done' }) }) + await expect(withoutEvidence.runChild({ + parentThreadId: 'thr_missing_evidence', + parentTurnId: 'turn_missing_evidence', + prompt: 'investigate', + returnFormat: 'evidence', + signal: new AbortController().signal + })).resolves.toMatchObject({ + status: 'failed', + error: 'child contract requires evidence but none was returned' + }) + }) + it('executes delegate_task through the normal tool host', async () => { const runtime = createRuntime() const host = new LocalToolHost({ @@ -109,6 +183,28 @@ describe('DelegationRuntime', () => { } }) + it('rejects invalid delegate_task budgets before starting a child', async () => { + const runtime = createRuntime() + const host = new LocalToolHost({ + registry: new CapabilityRegistry(buildDelegationToolProviders(runtime)) + }) + const result = await host.execute({ + callId: 'call_invalid_budget', + toolName: 'delegate_task', + arguments: { prompt: 'Investigate', tokenBudget: 0 } + }, { + threadId: 'thr_1', + turnId: 'turn_1', + workspace: '/tmp/ws', + approvalPolicy: 'auto', + abortSignal: new AbortController().signal, + awaitApproval: async () => 'allow' + }) + + expect(result.item).toMatchObject({ kind: 'tool_result', isError: true }) + expect((await runtime.diagnostics('thr_1')).childRuns).toEqual([]) + }) + it('caps concurrency at maxParallel and queues the overflow instead of erroring', async () => { const gate = deferred() let active = 0 diff --git a/kun/tests/http-server.test.ts b/kun/tests/http-server.test.ts index 3868a3b9d..a3273c3cb 100644 --- a/kun/tests/http-server.test.ts +++ b/kun/tests/http-server.test.ts @@ -170,6 +170,87 @@ describe('HTTP server', () => { expect(response.status).toBe(401) }) + it('reports and clears MCP OAuth diagnostics through the HTTP layer', async () => { + const h = buildHarness() + h.runtime.mcpOAuth = () => [ + { + serverId: 'google_drive', + enabled: true, + configured: true, + transport: 'streamable-http', + url: 'https://drivemcp.googleapis.com/mcp/v1', + status: 'authorized', + hasClientInformation: true, + hasTokens: true, + hasRefreshToken: true, + hasCodeVerifier: false, + hasDiscoveryState: true + } + ] + h.runtime.clearMcpOAuth = async (serverId?: string) => ({ cleared: serverId ? [serverId] : ['google_drive'] }) + + const listed = await dispatchRequest( + h.router, + new Request('http://localhost/v1/mcp/oauth', { + headers: { authorization: 'Bearer tok-1' } + }) + ) + expect(listed.status).toBe(200) + await expect(readJson(listed)).resolves.toMatchObject({ + servers: [{ serverId: 'google_drive', status: 'authorized', hasTokens: true }] + }) + + const cleared = await dispatchRequest( + h.router, + new Request('http://localhost/v1/mcp/oauth/google_drive', { + method: 'DELETE', + headers: { authorization: 'Bearer tok-1' } + }) + ) + expect(cleared.status).toBe(200) + await expect(readJson(cleared)).resolves.toEqual({ cleared: ['google_drive'] }) + }) + + it('runs MCP OAuth authorization through the HTTP layer', async () => { + const h = buildHarness() + const authorized: string[] = [] + h.runtime.authorizeMcpOAuth = async (serverId: string) => { + authorized.push(serverId) + return { serverId, status: 'authorized', authorized: true } + } + + const response = await dispatchRequest( + h.router, + new Request('http://localhost/v1/mcp/oauth/google_drive', { + method: 'POST', + headers: { authorization: 'Bearer tok-1' } + }) + ) + + expect(response.status).toBe(200) + await expect(readJson(response)).resolves.toEqual({ + serverId: 'google_drive', + status: 'authorized', + authorized: true + }) + expect(authorized).toEqual(['google_drive']) + }) + + it('reports MCP OAuth authorization as unavailable when the runtime lacks it', async () => { + const h = buildHarness() + h.runtime.authorizeMcpOAuth = undefined + + const response = await dispatchRequest( + h.router, + new Request('http://localhost/v1/mcp/oauth/google_drive', { + method: 'POST', + headers: { authorization: 'Bearer tok-1' } + }) + ) + + expect(response.status).toBe(503) + }) + it('lists discovered skills through the HTTP layer', async () => { const h = buildHarness() h.runtime.skills = () => ({ diff --git a/kun/tests/local-workspace-inspector.test.ts b/kun/tests/local-workspace-inspector.test.ts new file mode 100644 index 000000000..8f1eb961d --- /dev/null +++ b/kun/tests/local-workspace-inspector.test.ts @@ -0,0 +1,48 @@ +import type { ExecFileOptions } from 'node:child_process' +import { mkdir, mkdtemp, rm } from 'node:fs/promises' +import { tmpdir } from 'node:os' +import { join } from 'node:path' +import { describe, expect, it } from 'vitest' +import { LocalWorkspaceInspector } from '../src/adapters/workspace/local-workspace-inspector.js' + +describe('LocalWorkspaceInspector', () => { + it('runs git commands in the selected workspace', async () => { + const root = await mkdtemp(join(tmpdir(), 'kun-workspace-')) + const workspace = join(root, 'packages', 'app') + await mkdir(workspace, { recursive: true }) + + const calls: Array<{ args: string[]; cwd?: string }> = [] + const inspector = new LocalWorkspaceInspector({ + exec: async (_file: string, args: string[], options?: ExecFileOptions) => { + calls.push({ args, cwd: typeof options?.cwd === 'string' ? options.cwd : undefined }) + const command = args.join(' ') + if (command === 'rev-parse --is-inside-work-tree') return { stdout: 'true\n', stderr: '' } + if (command === 'rev-parse --abbrev-ref HEAD') return { stdout: 'develop\n', stderr: '' } + if (command === 'rev-parse HEAD') return { stdout: 'abc123\n', stderr: '' } + if (command === 'status --porcelain') return { stdout: ' M src/app.ts\n', stderr: '' } + throw new Error(`unexpected git command: ${command}`) + } + }) + + try { + const status = await inspector.status(workspace) + + expect(status).toMatchObject({ + path: workspace, + isGitRepository: true, + branch: 'develop', + headSha: 'abc123', + isDirty: true, + fileChangeCount: 1 + }) + expect(calls).toEqual([ + { args: ['rev-parse', '--is-inside-work-tree'], cwd: workspace }, + { args: ['rev-parse', '--abbrev-ref', 'HEAD'], cwd: workspace }, + { args: ['rev-parse', 'HEAD'], cwd: workspace }, + { args: ['status', '--porcelain'], cwd: workspace } + ]) + } finally { + await rm(root, { recursive: true, force: true }) + } + }) +}) diff --git a/kun/tests/loop.test.ts b/kun/tests/loop.test.ts index ffae9b5aa..d73293b51 100644 --- a/kun/tests/loop.test.ts +++ b/kun/tests/loop.test.ts @@ -67,6 +67,52 @@ describe('AgentLoop', () => { expect(request.tools.map((tool) => tool.name)).toContain('bash') expect(request.contextInstructions?.join('\n')).toContain('') expect(request.contextInstructions?.join('\n')).toContain('') + expect(request.contextInstructions?.join('\n')).not.toContain('Specialized MCP tools are available') + }) + + it('prefers specialized MCP tools only when they are advertised', async () => { + let observedRequest: ModelRequest | null = null + const sourceExplorer = LocalToolHost.defineTool({ + name: 'mcp_semantic_find_symbol', + description: 'Find source-code symbols and their references.', + inputSchema: { type: 'object' }, + policy: 'auto', + execute: async () => ({ output: {} }) + }) + const registry = new CapabilityRegistry([ + { + id: 'builtin', + kind: 'built-in', + enabled: true, + available: true, + tools: buildDefaultLocalTools() + }, + { + id: 'mcp:semantic', + kind: 'mcp', + enabled: true, + available: true, + tools: [sourceExplorer] + } + ]) + const h = makeHarness({ + provider: 'tool-preference', + model: 'tool-preference', + async *stream(request: ModelRequest): AsyncIterable { + observedRequest = request + yield { kind: 'completed', stopReason: 'stop' } + } + }, { toolHost: new LocalToolHost({ registry }) }) + await bootstrapThread(h) + + await h.loop.runTurn(h.threadId, h.turnId) + + const request = observedRequest as ModelRequest | null + if (!request) throw new Error('expected model request') + const instructions = request.contextInstructions?.join('\n') ?? '' + expect(instructions).toContain('Specialized source-code MCP tools are available') + expect(instructions).toContain('`mcp_semantic_find_symbol`') + expect(instructions).toContain('before broad `read`/`grep`/`find`/`ls` scans') }) it('records elapsed seconds for active goals after a turn finishes', async () => { diff --git a/kun/tests/mcp-config.test.ts b/kun/tests/mcp-config.test.ts index 1aaa00e9c..e38d4e86d 100644 --- a/kun/tests/mcp-config.test.ts +++ b/kun/tests/mcp-config.test.ts @@ -55,6 +55,32 @@ describe('MCP config', () => { expect(config.mcp.servers.github?.transport).toBe('streamable-http') }) + it('accepts OAuth settings for remote MCP servers', () => { + const server = McpServerConfig.parse({ + transport: 'streamable-http', + url: 'https://mcp.example.test/mcp', + trustScope: 'user', + oauth: { + clientName: 'Kun Test Client', + clientId: 'client-id', + clientSecret: 'client-secret', + scopes: ['drive.readonly', 'gmail.readonly'], + redirectPort: 49_999, + callbackTimeoutMs: 30_000 + } + }) + + expect(server.oauth).toMatchObject({ + enabled: true, + clientName: 'Kun Test Client', + clientId: 'client-id', + clientSecret: 'client-secret', + scopes: ['drive.readonly', 'gmail.readonly'], + redirectPort: 49_999, + callbackTimeoutMs: 30_000 + }) + }) + it('rejects stdio servers without commands', () => { const result = McpServerConfig.safeParse({ transport: 'stdio', diff --git a/kun/tests/mcp-tool-provider.test.ts b/kun/tests/mcp-tool-provider.test.ts index b7702854b..68483db6d 100644 --- a/kun/tests/mcp-tool-provider.test.ts +++ b/kun/tests/mcp-tool-provider.test.ts @@ -1,12 +1,22 @@ import { describe, expect, it } from 'vitest' +import { mkdtemp } from 'node:fs/promises' +import { createServer, get as httpGet } from 'node:http' +import type { AddressInfo } from 'node:net' +import { tmpdir } from 'node:os' +import { join } from 'node:path' import { CapabilityRegistry } from '../src/adapters/tool/capability-registry.js' import { LocalToolHost } from '../src/adapters/tool/local-tool-host.js' import { + FileMcpOAuthProvider, buildMcpStdioEnvironment, buildMcpToolProviders, + clearMcpOAuthCredentials, + createMcpOAuthProvider, formatMcpConnectionError, isMcpServerTrusted, isMcpServerVisible, + listMcpOAuthDiagnostics, + McpAuthorizationRequiredError, normalizeMcpToolName, resolveMcpServerCwd, type McpClientLike @@ -57,6 +67,30 @@ function fakeClient(): McpClientLike { } } +async function getFreePort(): Promise { + const server = createServer() + await new Promise((resolve, reject) => { + server.once('error', reject) + server.listen(0, '127.0.0.1', () => resolve()) + }) + const address = server.address() as AddressInfo + await new Promise((resolve, reject) => { + server.close((error) => error ? reject(error) : resolve()) + }) + return address.port +} + +async function httpStatus(url: URL): Promise { + return new Promise((resolve, reject) => { + const request = httpGet(url, (response) => { + response.resume() + response.on('end', () => resolve(response.statusCode ?? 0)) + }) + request.once('error', reject) + request.setTimeout(3_000, () => request.destroy(new Error('HTTP request timed out'))) + }) +} + describe('MCP tool provider', () => { it('normalizes stable MCP tool names', () => { expect(normalizeMcpToolName('GitHub Server', 'Search Issues')).toBe('mcp_github_server_search_issues') @@ -778,6 +812,232 @@ describe('MCP tool provider', () => { expect(encoded).not.toContain('config-secret') }) + it('keeps OAuth disabled unless remote MCP servers opt in', async () => { + const root = await mkdtemp(join(tmpdir(), 'kun-mcp-oauth-')) + const config = KunCapabilitiesConfig.parse({ + mcp: { + enabled: true, + servers: { + remote_docs: { + transport: 'streamable-http', + url: 'https://mcp.example.test/mcp', + trustScope: 'user' + } + } + } + }) + const server = config.mcp.servers.remote_docs as McpServerConfig + + expect(createMcpOAuthProvider('remote_docs', server, { storageDir: root })).toBeUndefined() + await expect(listMcpOAuthDiagnostics(config.mcp, { storageDir: root })).resolves.toEqual([]) + }) + + it('persists remote MCP OAuth client state outside the server config', async () => { + const root = await mkdtemp(join(tmpdir(), 'kun-mcp-oauth-')) + const server = KunCapabilitiesConfig.parse({ + mcp: { + enabled: true, + servers: { + google_drive: { + transport: 'streamable-http', + url: 'https://drivemcp.googleapis.com/mcp/v1', + trustScope: 'user', + oauth: { + scopes: ['drive.readonly'] + } + } + } + } + }).mcp.servers.google_drive as McpServerConfig + const storagePath = join(root, 'google_drive.json') + const provider = new FileMcpOAuthProvider('google_drive', server, storagePath, async () => undefined) + + await provider.saveClientInformation({ client_id: 'client-1', client_secret: 'secret-1' }) + await provider.saveTokens({ access_token: 'access-1', token_type: 'Bearer', refresh_token: 'refresh-1' }) + await provider.saveCodeVerifier('verifier-1') + + const restored = new FileMcpOAuthProvider('google_drive', server, storagePath, async () => undefined) + expect(await restored.clientInformation()).toMatchObject({ client_id: 'client-1' }) + expect(await restored.tokens()).toMatchObject({ access_token: 'access-1', refresh_token: 'refresh-1' }) + expect(await restored.codeVerifier()).toBe('verifier-1') + expect(restored.clientMetadata.scope).toBe('drive.readonly') + }) + + it('reports and clears remote MCP OAuth credential state', async () => { + const root = await mkdtemp(join(tmpdir(), 'kun-mcp-oauth-')) + const config = KunCapabilitiesConfig.parse({ + mcp: { + enabled: true, + servers: { + google_drive: { + transport: 'streamable-http', + url: 'https://drivemcp.googleapis.com/mcp/v1', + trustScope: 'user', + oauth: { + scopes: ['drive.readonly'] + } + } + } + } + }) + const server = config.mcp.servers.google_drive as McpServerConfig + const provider = createMcpOAuthProvider('google_drive', server, { storageDir: root }) + expect(provider).toBeDefined() + await provider?.saveTokens({ access_token: 'access-1', token_type: 'Bearer', refresh_token: 'refresh-1' }) + + const before = await listMcpOAuthDiagnostics(config.mcp, { storageDir: root }) + expect(before).toHaveLength(1) + expect(before[0]).toMatchObject({ + serverId: 'google_drive', + configured: true, + status: 'authorized', + hasTokens: true, + hasRefreshToken: true + }) + + const built = await buildMcpToolProviders(config.mcp, { + oauthStorageDir: root, + clientFactory: async () => fakeClient() + }) + await built.close() + expect(built.oauth[0]).toMatchObject({ + serverId: 'google_drive', + status: 'authorized' + }) + await expect(clearMcpOAuthCredentials(config.mcp, { + storageDir: root, + serverId: 'google_drive' + })).resolves.toEqual({ cleared: ['google_drive'] }) + expect((await listMcpOAuthDiagnostics(config.mcp, { storageDir: root }))[0]).toMatchObject({ + status: 'empty', + hasTokens: false + }) + }) + + it('receives remote MCP OAuth authorization codes on a loopback callback', async () => { + const root = await mkdtemp(join(tmpdir(), 'kun-mcp-oauth-')) + const redirectPort = await getFreePort() + const opened: string[] = [] + const server = KunCapabilitiesConfig.parse({ + mcp: { + enabled: true, + servers: { + vercel: { + transport: 'streamable-http', + url: 'https://mcp.vercel.com', + trustScope: 'user', + oauth: { + redirectPort, + callbackTimeoutMs: 5_000 + } + } + } + } + }).mcp.servers.vercel as McpServerConfig + const provider = new FileMcpOAuthProvider( + 'vercel', + server, + join(root, 'vercel.json'), + async (url) => { + opened.push(url.toString()) + }, + undefined, + true + ) + + const oauthState = provider.state() + await provider.redirectToAuthorization(new URL('https://auth.example.test/authorize')) + expect(provider.redirectUrl.port).toBe(String(redirectPort)) + const code = provider.waitForAuthorizationCode() + .then((value) => ({ ok: true as const, value })) + .catch((error: unknown) => ({ ok: false as const, error })) + const callbackUrl = new URL(provider.redirectUrl) + callbackUrl.searchParams.set('code', 'abc123') + callbackUrl.searchParams.set('state', oauthState) + const status = await httpStatus(callbackUrl) + + expect(status).toBe(200) + await expect(code).resolves.toEqual({ ok: true, value: 'abc123' }) + expect(opened).toEqual(['https://auth.example.test/authorize']) + }, 10_000) + + it('rejects non-http MCP OAuth authorization urls', async () => { + const root = await mkdtemp(join(tmpdir(), 'kun-mcp-oauth-')) + const opened: string[] = [] + const server = KunCapabilitiesConfig.parse({ + mcp: { + enabled: true, + servers: { + vercel: { + transport: 'streamable-http', + url: 'https://mcp.vercel.com', + trustScope: 'user', + oauth: {} + } + } + } + }).mcp.servers.vercel as McpServerConfig + const provider = new FileMcpOAuthProvider( + 'vercel', + server, + join(root, 'vercel.json'), + async (url) => { + opened.push(url.toString()) + }, + undefined, + true + ) + + await expect(provider.redirectToAuthorization(new URL('file:///tmp/token'))).rejects.toThrow(/http or https/) + expect(opened).toEqual([]) + }) + + it('keeps remote MCP OAuth callbacks bound to the generated state', async () => { + const root = await mkdtemp(join(tmpdir(), 'kun-mcp-oauth-')) + const redirectPort = await getFreePort() + const server = KunCapabilitiesConfig.parse({ + mcp: { + enabled: true, + servers: { + vercel: { + transport: 'streamable-http', + url: 'https://mcp.vercel.com', + trustScope: 'user', + oauth: { + redirectPort, + callbackTimeoutMs: 5_000 + } + } + } + } + }).mcp.servers.vercel as McpServerConfig + const provider = new FileMcpOAuthProvider( + 'vercel', + server, + join(root, 'vercel.json'), + async () => undefined, + undefined, + true + ) + + const state = provider.state() + await provider.redirectToAuthorization(new URL('https://auth.example.test/authorize')) + const code = provider.waitForAuthorizationCode() + .then((value) => ({ ok: true as const, value })) + .catch((error: unknown) => ({ ok: false as const, error })) + + const wrongStateUrl = new URL(provider.redirectUrl) + wrongStateUrl.searchParams.set('code', 'wrong-code') + wrongStateUrl.searchParams.set('state', 'wrong-state') + await expect(httpStatus(wrongStateUrl)).resolves.toBe(400) + + const callbackUrl = new URL(provider.redirectUrl) + callbackUrl.searchParams.set('code', 'right-code') + callbackUrl.searchParams.set('state', state) + await expect(httpStatus(callbackUrl)).resolves.toBe(200) + await expect(code).resolves.toEqual({ ok: true, value: 'right-code' }) + }, 10_000) + it('closes connected MCP clients during shutdown', async () => { let closed = 0 const config = KunCapabilitiesConfig.parse({ @@ -812,3 +1072,224 @@ describe('MCP tool provider', () => { expect(closed).toBe(1) }) }) + +describe('mcp oauth diagnostics state machine', () => { + function oauthServer(): McpServerConfig { + return KunCapabilitiesConfig.parse({ + mcp: { + enabled: true, + servers: { + vercel: { + transport: 'streamable-http', + url: 'https://mcp.vercel.com', + trustScope: 'user', + oauth: { scopes: ['projects.read'] } + } + } + } + }).mcp.servers.vercel as McpServerConfig + } + + it('reports empty then partial as credential material accrues', async () => { + const root = await mkdtemp(join(tmpdir(), 'kun-mcp-oauth-')) + const provider = new FileMcpOAuthProvider('vercel', oauthServer(), join(root, 'vercel.json'), async () => undefined) + + expect((await provider.diagnostics()).status).toBe('empty') + + await provider.saveCodeVerifier('verifier-1') + expect((await provider.diagnostics()).status).toBe('partial') + }) + + it('treats a saved token with future expiry as authorized and exposes expiresAt', async () => { + const root = await mkdtemp(join(tmpdir(), 'kun-mcp-oauth-')) + let clock = 1_700_000_000_000 + const provider = new FileMcpOAuthProvider( + 'vercel', + oauthServer(), + join(root, 'vercel.json'), + async () => undefined, + () => clock + ) + + await provider.saveTokens({ access_token: 'access-1', token_type: 'Bearer', expires_in: 3600 }) + const diagnostics = await provider.diagnostics() + + expect(diagnostics.status).toBe('authorized') + expect(diagnostics.expiresAt).toBe(new Date(clock + 3600 * 1000).toISOString()) + }) + + it('flips to expired once the access token outlives its lifetime', async () => { + const root = await mkdtemp(join(tmpdir(), 'kun-mcp-oauth-')) + let clock = 1_700_000_000_000 + const provider = new FileMcpOAuthProvider( + 'vercel', + oauthServer(), + join(root, 'vercel.json'), + async () => undefined, + () => clock + ) + + await provider.saveTokens({ access_token: 'access-1', token_type: 'Bearer', expires_in: 10, refresh_token: 'refresh-1' }) + clock += 20_000 + const diagnostics = await provider.diagnostics() + + expect(diagnostics.status).toBe('expired') + expect(diagnostics.hasRefreshToken).toBe(true) + }) + + it('surfaces the provider-granted scopes parsed from the token', async () => { + const root = await mkdtemp(join(tmpdir(), 'kun-mcp-oauth-')) + const provider = new FileMcpOAuthProvider('vercel', oauthServer(), join(root, 'vercel.json'), async () => undefined) + + await provider.saveTokens({ + access_token: 'access-1', + token_type: 'Bearer', + scope: 'projects.read deployments.read projects.read deployments.write' + }) + const diagnostics = await provider.diagnostics() + + expect(diagnostics.grantedScopes).toEqual(['projects.read', 'deployments.read', 'deployments.write']) + }) + + it('omits grantedScopes when the provider returns no scope', async () => { + const root = await mkdtemp(join(tmpdir(), 'kun-mcp-oauth-')) + const provider = new FileMcpOAuthProvider('vercel', oauthServer(), join(root, 'vercel.json'), async () => undefined) + + await provider.saveTokens({ access_token: 'access-1', token_type: 'Bearer' }) + const diagnostics = await provider.diagnostics() + + expect(diagnostics.grantedScopes).toBeUndefined() + }) + + it('surfaces a recorded authorization failure as error and clears it on the next token', async () => { + const root = await mkdtemp(join(tmpdir(), 'kun-mcp-oauth-')) + const provider = new FileMcpOAuthProvider('vercel', oauthServer(), join(root, 'vercel.json'), async () => undefined) + + await provider.recordAuthorizationError('MCP OAuth authorization failed: access_denied') + const failed = await provider.diagnostics() + expect(failed.status).toBe('error') + expect(failed.lastError).toContain('access_denied') + expect(failed.lastErrorAt).toBeDefined() + + await provider.saveTokens({ access_token: 'access-1', token_type: 'Bearer' }) + const recovered = await provider.diagnostics() + expect(recovered.status).toBe('authorized') + expect(recovered.lastError).toBeUndefined() + }) + + it('exposes an authorize entry point that no-ops for unknown servers', async () => { + const root = await mkdtemp(join(tmpdir(), 'kun-mcp-oauth-')) + const config = KunCapabilitiesConfig.parse({ + mcp: { enabled: true, servers: { vercel: oauthServer() } } + }) + const built = await buildMcpToolProviders(config.mcp, { + oauthStorageDir: root, + clientFactory: async () => fakeClient() + }) + await built.close() + + await expect(built.authorizeOAuth('does-not-exist')).resolves.toEqual({ + serverId: 'does-not-exist', + status: 'disabled', + authorized: false + }) + }) + + it('does not block startup when a remote server needs authorization', async () => { + const root = await mkdtemp(join(tmpdir(), 'kun-mcp-oauth-')) + const config = KunCapabilitiesConfig.parse({ + mcp: { enabled: true, servers: { vercel: oauthServer() } } + }) + let opened = 0 + // A non-interactive startup surfaces a typed "needs authorization" error + // instead of opening a browser; the build still resolves so the runtime is + // never blocked on a user completing an OAuth handshake. + const built = await buildMcpToolProviders(config.mcp, { + oauthStorageDir: root, + openExternal: () => { + opened += 1 + }, + clientFactory: async () => { + throw new McpAuthorizationRequiredError('vercel') + } + }) + await built.close() + + expect(opened).toBe(0) + const diagnostic = built.diagnostics.find((entry) => entry.id === 'vercel') + expect(diagnostic?.status).toBe('authorization_required') + expect(diagnostic?.lastError).toContain('Authorize') + }) + + it('refuses to open a browser from a non-interactive provider', async () => { + const root = await mkdtemp(join(tmpdir(), 'kun-mcp-oauth-')) + const opened: string[] = [] + const server = oauthServer() + const provider = new FileMcpOAuthProvider('vercel', server, join(root, 'vercel.json'), async (url) => { + opened.push(url.toString()) + }) + // Default (non-interactive): must throw before opening a browser/callback. + await expect(provider.redirectToAuthorization(new URL('https://auth.example.test/authorize'))) + .rejects.toThrow(/requires OAuth authorization/) + expect(opened).toEqual([]) + }) + + it('connects and registers a server immediately after a successful authorization', async () => { + const root = await mkdtemp(join(tmpdir(), 'kun-mcp-oauth-')) + const config = KunCapabilitiesConfig.parse({ + mcp: { enabled: true, servers: { vercel: oauthServer() } } + }) + const registered: string[] = [] + let authorizeCalls = 0 + let connectCalls = 0 + const built = await buildMcpToolProviders(config.mcp, { + oauthStorageDir: root, + // Startup: the server needs authorization, so the first connect fails and + // it must not connect/register yet. + clientFactory: async (serverId) => { + connectCalls += 1 + if (connectCalls === 1) throw new McpAuthorizationRequiredError(serverId) + return fakeClient() + }, + authorize: async (serverId) => { + authorizeCalls += 1 + return { serverId, status: 'authorized', authorized: true } + } + }) + // Capture the live register callback (as the runtime does). + await built.startBackgroundReconnect((provider) => registered.push(provider.id)) + + const result = await built.authorizeOAuth('vercel') + await built.close() + + expect(result).toMatchObject({ serverId: 'vercel', authorized: true }) + expect(authorizeCalls).toBe(1) + // The freshly authorized server is connected + registered live. + expect(registered).toContain('mcp:vercel') + expect(built.diagnostics.find((entry) => entry.id === 'vercel')?.status).toBe('connected') + }) + + it('shares one authorization run per server for concurrent clicks (single-flight)', async () => { + const root = await mkdtemp(join(tmpdir(), 'kun-mcp-oauth-')) + const config = KunCapabilitiesConfig.parse({ + mcp: { enabled: true, servers: { vercel: oauthServer() } } + }) + let authorizeCalls = 0 + const built = await buildMcpToolProviders(config.mcp, { + oauthStorageDir: root, + clientFactory: async () => fakeClient(), + authorize: async (serverId) => { + authorizeCalls += 1 + await new Promise((resolve) => setTimeout(resolve, 20)) + return { serverId, status: 'authorized', authorized: true } + } + }) + await built.startBackgroundReconnect(() => undefined) + + const [a, b] = await Promise.all([built.authorizeOAuth('vercel'), built.authorizeOAuth('vercel')]) + await built.close() + + expect(authorizeCalls).toBe(1) + expect(a).toEqual(b) + }) +}) diff --git a/kun/tests/memory-store.test.ts b/kun/tests/memory-store.test.ts index 9a2526212..a11d86439 100644 --- a/kun/tests/memory-store.test.ts +++ b/kun/tests/memory-store.test.ts @@ -1,4 +1,4 @@ -import { mkdtemp, readdir, readFile, rm } from 'node:fs/promises' +import { mkdir, mkdtemp, readdir, readFile, rm, writeFile } from 'node:fs/promises' import { tmpdir } from 'node:os' import { join, resolve } from 'node:path' import { afterEach, beforeEach, describe, expect, it } from 'vitest' @@ -6,7 +6,7 @@ import { CapabilityRegistry } from '../src/adapters/tool/capability-registry.js' import { LocalToolHost } from '../src/adapters/tool/local-tool-host.js' import { buildMemoryToolProviders } from '../src/adapters/tool/memory-tool-provider.js' import { KunCapabilitiesConfig, type MemoryCapabilityConfig } from '../src/contracts/capabilities.js' -import { FileMemoryStore } from '../src/memory/memory-store.js' +import { effectiveMemoryConfidence, FileMemoryStore } from '../src/memory/memory-store.js' import type { ModelClient, ModelRequest } from '../src/ports/model-client.js' import { dispatchRequest } from '../src/server/http-server.js' import { bootstrapThread, makeHarness } from './loop-test-harness.js' @@ -52,6 +52,77 @@ describe('Memory store and recall', () => { expect((await store.list({ workspace: '/tmp/ws', includeDeleted: true })).find((item) => item.id === memory.id)?.deletedAt).toBeTruthy() }) + it('tracks provenance, expiry, confidence decay, and legacy records', async () => { + let now = '2026-06-03T00:00:00.000Z' + const store = new FileMemoryStore({ + rootDir: join(dir, 'memory'), + config: memoryConfig(), + nowIso: () => now, + idGenerator: () => `mem_${nextId++}`, + confidenceHalfLifeMs: 1_000 + }) + const memory = await store.create({ + content: 'Observed build uses a generated manifest', + scope: 'workspace', + workspace: '/tmp/ws', + provenance: { kind: 'inference', turnId: 'turn_1', origin: 'analysis' }, + ttlMs: 2_000 + }) + expect(memory).toMatchObject({ + confidence: 0.4, + provenance: { kind: 'inference', turnId: 'turn_1', origin: 'analysis' }, + expiresAt: '2026-06-03T00:00:02.000Z' + }) + expect(effectiveMemoryConfidence(memory, Date.parse(now) + 1_000, 1_000)).toBeCloseTo(0.2) + expect(await store.retrieve({ query: 'generated manifest', workspace: '/tmp/ws', limit: 3 })).toHaveLength(1) + + now = '2026-06-03T00:00:03.000Z' + expect(await store.retrieve({ query: 'generated manifest', workspace: '/tmp/ws', limit: 3 })).toEqual([]) + expect((await store.diagnostics()).activeCount).toBe(0) + + await mkdir(join(dir, 'memory'), { recursive: true }) + await writeFile(join(dir, 'memory', 'legacy.json'), JSON.stringify({ + id: 'legacy', + content: 'Legacy user preference', + scope: 'user', + tags: [], + confidence: 1, + createdAt: now, + updatedAt: now + })) + expect((await store.list({ all: true })).find((item) => item.id === 'legacy')).toBeTruthy() + }) + + it('supersedes older memories and preserves user corrections', async () => { + const store = createStore() + const older = await store.create({ + content: 'Use npm for frontend installs', + scope: 'workspace', + workspace: '/tmp/ws' + }) + const newer = await store.create({ + content: 'Use pnpm for frontend installs', + scope: 'workspace', + workspace: '/tmp/ws', + provenance: { kind: 'tool', origin: 'package-manager-check' }, + supersedes: older.id + }) + + const records = await store.list({ workspace: '/tmp/ws' }) + expect(records.find((item) => item.id === older.id)?.supersededAt).toBeTruthy() + expect((await store.retrieve({ query: 'frontend installs', workspace: '/tmp/ws', limit: 3 })).map((item) => item.id)).toEqual([newer.id]) + + const corrected = await store.update(newer.id, { content: 'Use Bun for frontend installs' }, { + workspace: '/tmp/ws' + }) + expect(corrected).toMatchObject({ + content: 'Use Bun for frontend installs', + correctedFrom: 'Use pnpm for frontend installs', + confidence: 1, + provenance: { kind: 'user' } + }) + }) + it('exposes memory API routes with diagnostics', async () => { const h = buildHarness() h.runtime.memoryStore = createStore() @@ -153,7 +224,7 @@ describe('Memory store and recall', () => { await h.loop.runTurn(h.threadId, h.turnId) - expect(seenRequests.at(-1)?.contextInstructions?.[0]).toContain(memory.id) + expect(seenRequests.at(-1)?.contextInstructions?.join('\n')).toContain(memory.id) expect((await h.turns.getTurn(h.threadId, h.turnId))?.injectedMemoryIds).toEqual([memory.id]) expect((await store.diagnostics()).lastInjectedIds).toEqual([memory.id]) diff --git a/kun/tests/read-json-body.test.ts b/kun/tests/read-json-body.test.ts index 9ddf0071f..e37c4d920 100644 --- a/kun/tests/read-json-body.test.ts +++ b/kun/tests/read-json-body.test.ts @@ -35,4 +35,25 @@ describe('readJsonBody', () => { message: 'invalid JSON body' }) }) + + it('rejects a declared body that exceeds the configured byte limit', async () => { + const result = await readJsonBody(new Request('http://localhost/v1/demo', { + method: 'POST', + headers: { 'content-length': '128' }, + body: '{}' + }), 32) + expect(result.ok).toBe(false) + if (result.ok) return + expect(result.response.status).toBe(413) + }) + + it('rejects a streamed body that exceeds the configured byte limit', async () => { + const result = await readJsonBody(new Request('http://localhost/v1/demo', { + method: 'POST', + body: JSON.stringify({ text: 'x'.repeat(128) }) + }), 32) + expect(result.ok).toBe(false) + if (result.ok) return + expect(result.response.status).toBe(413) + }) }) diff --git a/release/release-v0.2.21.md b/release/release-v0.2.21.md new file mode 100644 index 000000000..b6bcc06e6 --- /dev/null +++ b/release/release-v0.2.21.md @@ -0,0 +1,69 @@ +# Kun v0.2.21 + +这一版是 v0.2.20 之后的一次 Agent 工具链与运行时治理增强。主线是让 Kun 更会“看项目、管上下文、接外部工具”：新增预算受控的代码地图、超大工具输出 artifact、文档附件文本注入、远程 MCP OAuth、任务图和更严格的子代理执行契约；同时修复 Git 检查点、分支导航、SSE 断线恢复、附件安全和多语言流式渲染等问题。 + +需要特别说明的是：此前进入 `develop` 的 SSH Remote / 远程运行位置 / 端口转发功能已在本版发布前撤回。这个方向后续会重新设计，不会作为 v0.2.21 的正式功能发布。 + +### 项目理解与工具输出 + +- 新增 `repo_map` 代码地图工具，可以在预算内扫描本地项目结构，帮助 Agent 更快建立代码上下文。 +- 超大工具输出会被转存为 artifact，避免把长输出直接塞满对话上下文；同时新增读取 artifact 的能力，后续可以按需查看完整内容或指定行范围。 +- MCP 工具检索进一步收敛：模型可以优先通过专门的 MCP inspection/search 流程定位工具，而不是每轮携带完整工具目录。 +- Headless `kun exec` 支持 JSONL 流式输出，方便脚本、CI 或外部集成实时消费运行状态。 + +### MCP、Marketplace 与安全 + +- 远程 HTTP/SSE MCP server 支持 OAuth 授权，Kun 会把 token 存在运行时数据目录，而不是写回 MCP 配置文件。 +- 新增 MCP OAuth 状态查看与清除接口，便于排查和撤销远程 MCP 授权。 +- Marketplace 安装会校验并固定 MCP package，减少远程包被替换或配置漂移带来的风险。 +- MCP stdio 环境、命名与传输层做了进一步拆分和加固，降低外部工具配置对运行时稳定性的影响。 +- 文档附件和外部内容会用 untrusted content 包裹后注入模型上下文，减少把用户文件内容误当系统指令的风险。 + +### 附件、文档与写作 + +- 文档附件现在会提取文本并注入到模型请求中，模型可以直接理解上传文档内容，而不只是看到文件名。 +- 附件上传边界进一步收紧，避免不安全路径或超出预期的本地文件被误传入运行时。 +- Write 助手可以看到运行时 Skill 能力，让写作工作台里的 Agent 更容易复用已有项目技能。 +- 修复越南语等组合字符在流式输出中被拆坏的问题，长文本渲染更可靠。 + +### 任务图、子代理与记忆治理 + +- 新增持久化 Agent task graph，为后续更复杂的多步骤任务调度和恢复打基础。 +- 子代理执行契约更严格，会约束只读/可写工具、输入输出和执行边界，减少委派任务越权或跑偏。 +- 长期记忆增加来源、有效期和治理信息，方便判断记忆从哪里来、是否还该继续使用。 +- 缓存命中率新增回归检测，usage 展示里的缓存命中率计算也更一致。 +- Replay benchmark 增加质量断言，让运行时回放不仅能跑，还能检查结果是否退化。 + +### Git、工作区与对话体验 + +- Git 分支前缀现在可配置，团队可以统一新分支命名风格。 +- Git 检查会在当前选中的工作区中运行，避免多项目场景下读错目录。 +- 默认开启未使用 Git checkpoint 清理，并修复部分 checkpoint restore 可能造成数据不完整的问题。 +- 改善分支对话导航，切换分支/会话时更不容易迷路。 +- 设置中的日志路径背景更清楚，深浅色界面下都更易读。 +- SDD 需求工作区会显式绑定，减少需求草稿写到错误 workspace 的可能。 + +### 运行时稳定性与回归修复 + +- 修复 SSE 断线后可能丢失 assistant 回复的问题，断线重连后的对话恢复更稳。 +- 修复 OpenCode Go DeepSeek v4 相关模型预设。 +- 修复 release validation 中暴露的测试期望、固定端口冲突和错误详情折叠问题，让 `develop` 的发版检查更稳定。 +- 补充附件、任务图、MCP OAuth、artifact、repo map、子代理、memory、usage 与 Git checkpoint 等测试覆盖。 + +### 已撤回的内容 + +- SSH Remote runtime tools、目标选择 UI 和 SSH port forwarding 在本版发布前已从 `develop` 移除。 +- 撤回原因是功能成熟度和用户体验设计还不够好；后续会重新设计远程执行体验，再以更完整的方式进入版本线。 +- 远程 MCP OAuth 不受影响：它是 MCP server 授权能力，不是 SSH Remote 执行功能。 + +### 升级说明 + +- 从 `v0.2.20` 升级可直接通过 GUI 更新。 +- 如果你使用远程 HTTP/SSE MCP server，可以在 MCP 配置中启用 OAuth；授权 token 会保存在 Kun runtime data 目录。 +- Marketplace 现在会更严格地校验 MCP package 和配置，旧的非固定或不安全配置可能需要重新确认后再安装。 +- 文档附件会把提取出的文本注入上下文；对于很大的文档，Kun 会做截断和安全包装。 +- SSH Remote 不在本版中发布，如果你在开发版里试用过该入口，升级到 v0.2.21 后将看不到相关运行位置和端口转发功能。 + +### 总结 + +v0.2.21 不是一个追求“堆大功能”的版本，而是把 Kun 的 Agent 底座继续做厚：更会读项目、更会管理长输出和文档、更安全地连接 MCP、更严格地约束子代理和记忆。同时，它也把尚未成熟的 SSH Remote 从发版线上撤下，给后续重新设计留出空间。 diff --git a/src/main/index.ts b/src/main/index.ts index 4e8238176..093da8211 100644 --- a/src/main/index.ts +++ b/src/main/index.ts @@ -72,6 +72,7 @@ import { waitForKunStartupSettled, type KunUnexpectedExitInfo } from './kun-process' +import { expandHomePath } from './settings-store' import { RestartBudget, type KunRuntimeStatus } from './kun-runtime-supervisor' import { configureLogger, logError, logWarn, pruneOnStartup } from './logger' import { cleanupUnusedGitCheckpointsIfDue } from './services/git-checkpoint-service' @@ -249,8 +250,15 @@ async function runCheckpointCleanupIfDue(settings: AppSettingsV1): Promise const runtime = resolveKunRuntimeSettings(settings) const dataDir = resolveKunDataDir(runtime) const intervalDays = settings.checkpointCleanup.intervalDays + const checkpointsRoot = settings.checkpointCleanup.directory?.trim() + ? expandHomePath(settings.checkpointCleanup.directory.trim()) + : undefined try { - const cleanup = await cleanupUnusedGitCheckpointsIfDue({ dataDir, intervalDays }) + const cleanup = await cleanupUnusedGitCheckpointsIfDue({ + dataDir, + intervalDays, + ...(checkpointsRoot ? { checkpointsRoot } : {}) + }) if (!cleanup.due) return const { result } = cleanup console.info( diff --git a/src/main/ipc/app-ipc-schemas.test.ts b/src/main/ipc/app-ipc-schemas.test.ts index 1828451e8..22f4fd927 100644 --- a/src/main/ipc/app-ipc-schemas.test.ts +++ b/src/main/ipc/app-ipc-schemas.test.ts @@ -47,6 +47,17 @@ describe('app-ipc-schemas', () => { expect(payload.path).toBe('/v1/runtime/tools') }) + it('accepts Kun MCP OAuth status and token reset endpoints', () => { + expect(runtimeRequestPayloadSchema.parse({ + path: '/v1/mcp/oauth', + method: 'GET' + }).path).toBe('/v1/mcp/oauth') + expect(runtimeRequestPayloadSchema.parse({ + path: '/v1/mcp/oauth/google_drive', + method: 'DELETE' + }).path).toBe('/v1/mcp/oauth/google_drive') + }) + it('accepts the Kun skills endpoint', () => { const payload = runtimeRequestPayloadSchema.parse({ path: '/v1/skills', diff --git a/src/main/ipc/app-ipc-schemas.ts b/src/main/ipc/app-ipc-schemas.ts index e1407d4cf..c9fbd0326 100644 --- a/src/main/ipc/app-ipc-schemas.ts +++ b/src/main/ipc/app-ipc-schemas.ts @@ -9,6 +9,8 @@ import { KUN_MEMORY_DIAGNOSTICS_TEMPLATE, KUN_MEMORY_RECORD_TEMPLATE, KUN_MEMORY_TEMPLATE, + KUN_MCP_OAUTH_SERVER_TEMPLATE, + KUN_MCP_OAUTH_TEMPLATE, KUN_RUNTIME_INFO_TEMPLATE, KUN_RUNTIME_TOOLS_TEMPLATE, KUN_SESSION_RESUME_TEMPLATE, @@ -160,6 +162,8 @@ const ENDPOINTS: readonly EndpointTemplate[] = [ compileEndpoint(KUN_MEMORY_TEMPLATE, ['GET', 'POST']), compileEndpoint(KUN_MEMORY_DIAGNOSTICS_TEMPLATE, ['GET']), compileEndpoint(KUN_MEMORY_RECORD_TEMPLATE, ['PATCH', 'DELETE']), + compileEndpoint(KUN_MCP_OAUTH_TEMPLATE, ['GET', 'DELETE']), + compileEndpoint(KUN_MCP_OAUTH_SERVER_TEMPLATE, ['DELETE']), compileEndpoint(KUN_THREADS_TEMPLATE, ['GET', 'POST']), compileEndpoint(KUN_THREAD_TEMPLATE, ['GET', 'PATCH', 'DELETE']), compileEndpoint(KUN_THREAD_FORK_TEMPLATE, ['POST']), @@ -534,7 +538,11 @@ const checkpointCleanupPatchSchema = z.object({ z.literal(3), z.literal(5), z.literal(10) - ]).optional() + ]).optional(), + // Issue #651: user-configurable checkpoint storage directory (e.g. another + // drive) + per-thread retention cap. Empty string clears the override. + directory: z.string().max(4096).optional(), + maxPerThread: z.number().int().min(1).max(100).optional() }).strict() const notificationsPatchSchema = z.object({ @@ -1364,6 +1372,7 @@ const settingsPatchObjectSchema = z.object({ conversationWorkspaceRoot: defaultPathSchema, log: logPatchSchema.optional(), checkpointCleanup: checkpointCleanupPatchSchema.optional(), + gitBranchPrefix: trimmedString(128).or(z.literal('')).optional(), notifications: notificationsPatchSchema.optional(), appBehavior: appBehaviorPatchSchema.optional(), keyboardShortcuts: keyboardShortcutsPatchSchema.optional(), @@ -1440,7 +1449,8 @@ export const gitCheckpointCreatePayloadSchema = z export const gitCheckpointRestorePayloadSchema = z .object({ - checkpointId: trimmedString(MAX_ID_LENGTH * 4) + checkpointId: trimmedString(MAX_ID_LENGTH * 4), + allowPartialRestore: z.boolean().optional() }) .strict() diff --git a/src/main/ipc/register-app-ipc-handlers.ts b/src/main/ipc/register-app-ipc-handlers.ts index c865def1a..415dfba94 100644 --- a/src/main/ipc/register-app-ipc-handlers.ts +++ b/src/main/ipc/register-app-ipc-handlers.ts @@ -127,7 +127,7 @@ import { removeGitBranchWorktree, switchGitBranch } from '../services/git-service' -import { createGitCheckpoint, restoreGitCheckpoint } from '../services/git-checkpoint-service' +import { createGitCheckpoint, restoreGitCheckpoint, type GitCheckpointStorageOptions } from '../services/git-checkpoint-service' import { abortMerge, abortRebase, @@ -1011,6 +1011,16 @@ export function registerAppIpcHandlers(options: RegisterAppIpcHandlersOptions): return expandHomePath(runtime.dataDir?.trim() || DEFAULT_KUN_DATA_DIR) } + // Map the user's checkpoint settings (issue #651) to the service storage + // options: an optional directory override (e.g. another drive) and the + // per-thread retention cap. Home-relative paths are expanded. + const resolveCheckpointStorageOptions = ( + cfg: { directory?: string; maxPerThread?: number } + ): GitCheckpointStorageOptions => ({ + ...(cfg.directory?.trim() ? { checkpointsRoot: expandHomePath(cfg.directory.trim()) } : {}), + ...(cfg.maxPerThread !== undefined ? { maxPerThread: cfg.maxPerThread } : {}) + }) + ipcMain.handle('kun:sessions:detect-legacy', async () => detectLegacySessions({ homeDir: homedir(), destDataDir: await resolveKunThreadsDataDir() }) ) @@ -1071,17 +1081,22 @@ export function registerAppIpcHandlers(options: RegisterAppIpcHandlersOptions): ) ipcMain.handle('git:checkpoint:create', async (_, payload: unknown) => { const request = parseIpcPayload('git:checkpoint:create', gitCheckpointCreatePayloadSchema, payload) + const settings = await store.load() return createGitCheckpoint({ dataDir: await resolveKunThreadsDataDir(), workspaceRoot: request.workspaceRoot, - threadId: request.threadId + threadId: request.threadId, + storage: resolveCheckpointStorageOptions(settings.checkpointCleanup) }) }) ipcMain.handle('git:checkpoint:restore', async (_, payload: unknown) => { const request = parseIpcPayload('git:checkpoint:restore', gitCheckpointRestorePayloadSchema, payload) + const settings = await store.load() return restoreGitCheckpoint({ dataDir: await resolveKunThreadsDataDir(), checkpointId: request.checkpointId, + ...(request.allowPartialRestore ? { allowPartialRestore: true } : {}), + storage: resolveCheckpointStorageOptions(settings.checkpointCleanup), // Bridge the main-process runtimeRequest into the shape restoreGitCheckpoint // expects ((path, {method, body}) => {ok,status,body}). On a transport-level // failure (runtime not up, connection refused) we return a non-ok result so diff --git a/src/main/kun-process.test.ts b/src/main/kun-process.test.ts index b902c4b45..9b9ffcd93 100644 --- a/src/main/kun-process.test.ts +++ b/src/main/kun-process.test.ts @@ -27,6 +27,7 @@ vi.mock('electron', () => ({ })) let tempRoot: string | null = null +let testKunPort = 18899 function createSettings(binaryPath: string): AppSettingsV1 { return { @@ -38,7 +39,7 @@ function createSettings(binaryPath: string): AppSettingsV1 { provider: defaultModelProviderSettings(), agents: { kun: { - ...defaultKunRuntimeSettings(18899), + ...defaultKunRuntimeSettings(testKunPort), binaryPath, autoStart: true } @@ -96,8 +97,24 @@ function canBindTestPort(port: number): Promise { }) } -beforeEach(() => { +function allocateTestPort(): Promise { + return new Promise((resolve, reject) => { + const server = createServer() + server.unref() + server.once('error', reject) + server.listen(0, '127.0.0.1', () => { + const address = server.address() + server.close(() => { + if (address && typeof address === 'object') resolve(address.port) + else reject(new Error('failed to allocate a test port')) + }) + }) + }) +} + +beforeEach(async () => { tempRoot = mkdtempSync(join(tmpdir(), 'kun-process-')) + testKunPort = await allocateTestPort() configureLogger({ dir: tempRoot, enabled: true, retentionDays: 7 }) }) @@ -117,7 +134,7 @@ describe('startKunChild', () => { 'ready-child.js', [ "const http = require('node:http')", - "const port = 18899", + `const port = ${testKunPort}`, "const server = http.createServer((req, res) => {", " res.setHeader('content-type', 'application/json')", " res.end(JSON.stringify({ service: 'kun', mode: 'serve', status: 'ok' }))", @@ -136,7 +153,7 @@ describe('startKunChild', () => { await module.stopKunChildAndWait() const logText = await readKunLog() expect(logText).toContain('KUN_READY') - expect(logText).toContain('ready marker received on port 18899') + expect(logText).toContain(`ready marker received on port ${testKunPort}`) }) it('does not settle on the ready marker until the /health endpoint responds', async () => { @@ -148,7 +165,7 @@ describe('startKunChild', () => { "const http = require('node:http')", "const { existsSync } = require('node:fs')", `const healthSignalPath = ${JSON.stringify(healthSignalPath)}`, - "const port = 18899", + `const port = ${testKunPort}`, // Emit the ready marker right away but serve no /health yet: the // marker alone must NOT be enough to settle the launch. "process.stdout.write('KUN_READY ' + JSON.stringify({ service: 'kun', mode: 'serve', port }) + '\\n')", @@ -195,7 +212,7 @@ describe('startKunChild', () => { "const http = require('node:http')", "const { existsSync } = require('node:fs')", `const readySignalPath = ${JSON.stringify(readySignalPath)}`, - "const port = 18899", + `const port = ${testKunPort}`, 'let sentReady = false', // Only stand up the /health server once the signal exists so the // parallel health probe cannot settle the launch before then. @@ -306,7 +323,7 @@ describe('waitForKunStartupSettled', () => { "const http = require('node:http')", "const { existsSync } = require('node:fs')", `const readySignalPath = ${JSON.stringify(readySignalPath)}`, - "const port = 18899", + `const port = ${testKunPort}`, 'let sentReady = false', // Only stand up the /health server once the signal exists so the // parallel health probe cannot settle the launch before then. @@ -841,7 +858,9 @@ describe('syncGuiManagedKunConfig', () => { expect(parsed.capabilities.skills.enabled).toBe(true) expect(parsed.capabilities.skills.legacySkillMd).toBe(true) expect(parsed.capabilities.skills.roots).toEqual(expect.arrayContaining([ - join(workspaceRoot, '.codex', 'skills'), + join(workspaceRoot, '.codex', 'skills') + ])) + expect(parsed.capabilities.skills.globalRoots).toEqual(expect.arrayContaining([ extraRoot ])) }) diff --git a/src/main/kun-process.ts b/src/main/kun-process.ts index f83751c43..9ee03c65d 100644 --- a/src/main/kun-process.ts +++ b/src/main/kun-process.ts @@ -643,9 +643,20 @@ async function skillCapabilityConfigForRuntime( path.length > 0 && !managed.has(comparableSkillRootPath(path)) && !isCodexPluginCacheRoot(path)) + const manualGlobalExisting = stringArrayValue(existing.globalRoots) + .map(normalizeSkillRootPath) + .filter((path) => + path.length > 0 && + !managed.has(comparableSkillRootPath(path)) && + !isCodexPluginCacheRoot(path)) + const guiRoots = await guiSkillRootsForRuntime(settings) const roots = uniqueStrings([ ...manualExisting, - ...(await guiSkillRootsForRuntime(settings)).map((root) => root.path) + ...guiRoots.filter((root) => root.scope === 'project').map((root) => root.path) + ]) + const globalRoots = uniqueStrings([ + ...manualGlobalExisting, + ...guiRoots.filter((root) => root.scope === 'global').map((root) => root.path) ]) return { ...existing, @@ -653,11 +664,11 @@ async function skillCapabilityConfigForRuntime( // enable toggle, so a persisted `enabled: false` is only ever the schema // default leaking onto disk — it must not permanently suppress discovered // skills. An explicit `true` still forces on even with no roots. - enabled: roots.length > 0 || existing.enabled === true, + enabled: roots.length > 0 || globalRoots.length > 0 || existing.enabled === true, roots, workspaceRoots: guiSkillWorkspaceRootsForRuntime(settings), // #149: Pass global skill roots from settings (e.g. ~/.kun/skills) - globalRoots: existing.globalRoots ?? [], + globalRoots, // Skills the user disabled in the GUI. Forwarded so the runtime drops them // from discovery — without this they stay loadable via load_skill and keep // appearing in the catalog despite the GUI toggle (#392). @@ -715,6 +726,7 @@ function normalizeGuiManagedMcpServer(server: unknown): Record const args = stringArrayValue(raw.args) const headers = stringRecordValue(raw.headers) const env = stringRecordValue(raw.env) + const oauth = objectValue(raw.oauth) const transport = normalizeMcpTransport(raw.transport, command, url) if (!transport) return null @@ -734,6 +746,7 @@ function normalizeGuiManagedMcpServer(server: unknown): Record ...(Object.keys(headers).length > 0 ? { headers } : {}), ...(Object.keys(env).length > 0 ? { env } : {}), ...(workspaceRoots.length > 0 ? { workspaceRoots } : {}), + ...(Object.keys(oauth).length > 0 ? { oauth } : {}), trustScope, ...(trustedWorkspaceRoots.length > 0 ? { trustedWorkspaceRoots } : {}), ...(timeoutMs ? { timeoutMs } : {}) diff --git a/src/main/runtime-sse-ipc.test.ts b/src/main/runtime-sse-ipc.test.ts index 572088c6e..051a91d5e 100644 --- a/src/main/runtime-sse-ipc.test.ts +++ b/src/main/runtime-sse-ipc.test.ts @@ -64,6 +64,9 @@ describe('runtime-sse-ipc', () => { if (chunk === '__ERROR__') { throw new Error('Network Disruption') } + if (chunk === '__TERMINATED__') { + throw new Error('terminated') + } return { done: false, value: enc.encode(chunk) } } } @@ -167,4 +170,63 @@ describe('runtime-sse-ipc', () => { expect(allEvents[2].seq).toBe(3) expect(allEvents[2].text).toBe('bye') }) + + it('treats terminated stream reads as reconnectable SSE disconnects', async () => { + registerRuntimeSseIpc({ + ipcMain: mockIpcMain, + store: mockStore, + ensureRuntime: mockEnsureRuntime, + logError: mockLogError + }) + + const startHandler = handlers.get('runtime:sse:start') + expect(startHandler).toBeDefined() + + const stream1 = mockReadableStream([ + 'id: 7\ndata: {"text": "partial"}\n\n', + '__TERMINATED__' + ]) + const stream2 = mockReadableStream([ + 'id: 8\ndata: {"text": "final"}\n\n' + ]) + + mockFetch.mockImplementation(async () => { + const callCount = mockFetch.mock.calls.length + if (callCount === 1) { + return { ok: true, status: 200, body: stream1 } + } + if (callCount === 2) { + return { ok: true, status: 200, body: stream2 } + } + return { ok: false, status: 400 } + }) + + const startRes = await startHandler!(mockEvent, { + threadId: 'thread-terminated', + sinceSeq: 0 + }) + + await vi.advanceTimersByTimeAsync(750) + + expect(mockFetch).toHaveBeenCalledTimes(3) + expect(mockFetch.mock.calls[1][0].toString()).toContain('since_seq=7') + expect(mockEvent.sender.send).not.toHaveBeenCalledWith( + 'runtime:sse-error', + expect.objectContaining({ streamId: startRes.streamId, message: 'terminated' }) + ) + expect(mockLogError).not.toHaveBeenCalledWith( + 'sse', + expect.stringContaining('SSE stream error'), + expect.objectContaining({ message: 'terminated' }) + ) + + const stopHandler = handlers.get('runtime:sse:stop') + expect(stopHandler).toBeDefined() + await stopHandler!(mockEvent, startRes.streamId) + + const allEvents = mockEvent.sender.send.mock.calls + .filter((call: any) => call[0] === 'runtime:sse-event') + .flatMap((call: any) => call[1].events) + expect(allEvents.map((event: any) => event.seq)).toEqual([7, 8]) + }) }) diff --git a/src/main/runtime-sse-ipc.ts b/src/main/runtime-sse-ipc.ts index 3af3e7db2..152747c7d 100644 --- a/src/main/runtime-sse-ipc.ts +++ b/src/main/runtime-sse-ipc.ts @@ -111,6 +111,10 @@ function isFatalSseStatus(status: number | undefined): boolean { return typeof status === 'number' && status >= 400 && status < 500 && status !== 408 && status !== 429 } +function isTransientSseErrorMessage(message: string): boolean { + return /sse start timeout|fetch failed|network|terminated|aborted|socket|ECONNRESET|ECONNREFUSED|ETIMEDOUT|EPIPE|UND_ERR/i.test(message) +} + async function fetchSseWithStartTimeout( url: URL, headers: Record, @@ -285,7 +289,7 @@ export function registerRuntimeSseIpc(options: { } catch (e) { if (state.stoppedByClient || ac.signal.aborted) return const msg = e instanceof Error ? e.message : String(e) - if (/sse start timeout/i.test(msg) || /fetch failed/i.test(msg) || /network/i.test(msg)) { + if (isTransientSseErrorMessage(msg)) { await sleepWithAbort(reconnectDelayMs, ac.signal) reconnectDelayMs = Math.min(reconnectDelayMs * 2, SSE_RECONNECT_MAX_MS) continue diff --git a/src/main/services/git-checkpoint-service.test.ts b/src/main/services/git-checkpoint-service.test.ts index 8d60eec95..10742f79f 100644 --- a/src/main/services/git-checkpoint-service.test.ts +++ b/src/main/services/git-checkpoint-service.test.ts @@ -417,3 +417,160 @@ describe('git checkpoint service', () => { expect(second.result.deletedIds).toEqual(['gcp_second']) }) }) + +describe('git checkpoint storage limits (issue #651)', () => { + it('stores checkpoints under a user-configured directory (e.g. another drive)', async () => { + const customRoot = join(sandbox, 'other-drive', 'kun-checkpoints') + const checkpoint = await createGitCheckpoint({ + dataDir, + workspaceRoot: repoRoot, + threadId: 'thr_1', + storage: { checkpointsRoot: customRoot } + }) + expect(checkpoint.ok).toBe(true) + if (!checkpoint.ok) throw new Error(checkpoint.message) + await expect(stat(join(customRoot, checkpoint.checkpointId, 'metadata.json'))).resolves.toBeTruthy() + // Nothing should have been written under the default data dir location. + await expect(stat(join(dataDir, 'git-checkpoints', checkpoint.checkpointId))).rejects.toBeTruthy() + }) + + it('skips untracked files larger than the per-file cap and records them', async () => { + await writeFile(join(repoRoot, 'small.txt'), 'tiny') + await writeFile(join(repoRoot, 'huge.bin'), Buffer.alloc(2_000_000, 1)) + const checkpoint = await createGitCheckpoint({ + dataDir, + workspaceRoot: repoRoot, + threadId: 'thr_1', + storage: { maxUntrackedFileBytes: 1_000_000 } + }) + expect(checkpoint.ok).toBe(true) + if (!checkpoint.ok) throw new Error(checkpoint.message) + const dir = join(dataDir, 'git-checkpoints', checkpoint.checkpointId) + const metadata = JSON.parse(await readFile(join(dir, 'metadata.json'), 'utf-8')) as { + untrackedFiles: string[]; skippedUntracked?: string[] + } + expect(metadata.untrackedFiles).toContain('small.txt') + expect(metadata.skippedUntracked).toContain('huge.bin') + await expect(stat(join(dir, 'untracked', 'huge.bin'))).rejects.toBeTruthy() + await expect(stat(join(dir, 'untracked', 'small.txt'))).resolves.toBeTruthy() + }) + + it('stops snapshotting untracked files once the total budget is hit', async () => { + await writeFile(join(repoRoot, 'a.bin'), Buffer.alloc(600_000, 1)) + await writeFile(join(repoRoot, 'b.bin'), Buffer.alloc(600_000, 1)) + const checkpoint = await createGitCheckpoint({ + dataDir, + workspaceRoot: repoRoot, + threadId: 'thr_1', + storage: { maxUntrackedFileBytes: 1_000_000, maxUntrackedTotalBytes: 1_000_000 } + }) + if (!checkpoint.ok) throw new Error(checkpoint.message) + const dir = join(dataDir, 'git-checkpoints', checkpoint.checkpointId) + const metadata = JSON.parse(await readFile(join(dir, 'metadata.json'), 'utf-8')) as { + untrackedFiles: string[]; skippedUntracked?: string[] + } + // One file fits the 1MB budget, the second is skipped. + expect(metadata.untrackedFiles.length).toBe(1) + expect(metadata.skippedUntracked?.length).toBe(1) + }) + + it('marks a checkpoint with skipped untracked files as partial and refuses to restore it (no data loss)', async () => { + // A large untracked file is skipped by the size cap, so the checkpoint is + // partial. Restoring would `git clean -fd` the never-captured file, so the + // restore must be refused unless the caller opts in. + await writeFile(join(repoRoot, 'huge.bin'), Buffer.alloc(2_000_000, 1)) + const checkpoint = await createGitCheckpoint({ + dataDir, + workspaceRoot: repoRoot, + threadId: 'thr_partial', + storage: { maxUntrackedFileBytes: 1_000_000 } + }) + if (!checkpoint.ok) throw new Error(checkpoint.message) + const dir = join(dataDir, 'git-checkpoints', checkpoint.checkpointId) + const metadata = JSON.parse(await readFile(join(dir, 'metadata.json'), 'utf-8')) as { completeness?: string } + expect(metadata.completeness).toBe('partial') + + const restored = await restoreGitCheckpoint({ dataDir, checkpointId: checkpoint.checkpointId }) + expect(restored.ok).toBe(false) + if (restored.ok) throw new Error('expected partial restore to be refused') + expect(restored.reason).toBe('partial') + expect('skippedUntracked' in restored && restored.skippedUntracked).toContain('huge.bin') + // The destructive ops never ran: the skipped file is byte-for-byte intact. + expect((await stat(join(repoRoot, 'huge.bin'))).size).toBe(2_000_000) + }) + + it('marks a fully-captured checkpoint as complete', async () => { + await writeFile(join(repoRoot, 'small.txt'), 'tiny') + const checkpoint = await createGitCheckpoint({ dataDir, workspaceRoot: repoRoot, threadId: 'thr_complete' }) + if (!checkpoint.ok) throw new Error(checkpoint.message) + const dir = join(dataDir, 'git-checkpoints', checkpoint.checkpointId) + const metadata = JSON.parse(await readFile(join(dir, 'metadata.json'), 'utf-8')) as { completeness?: string } + expect(metadata.completeness).toBe('complete') + }) + + it('restores a partial checkpoint only when the bounded rescue is complete', async () => { + await writeFile(join(repoRoot, 'huge.bin'), Buffer.alloc(2_000_000, 7)) + const checkpoint = await createGitCheckpoint({ + dataDir, + workspaceRoot: repoRoot, + threadId: 'thr_partial_ok', + storage: { maxUntrackedFileBytes: 1_000_000 } + }) + if (!checkpoint.ok) throw new Error(checkpoint.message) + + const restored = await restoreGitCheckpoint({ + dataDir, + checkpointId: checkpoint.checkpointId, + allowPartialRestore: true + }) + expect(restored.ok).toBe(true) + if (!restored.ok) throw new Error(restored.message) + expect(restored.rescueCheckpointId).toMatch(/^gcp_/) + // The file exceeds the original checkpoint's custom cap but fits the normal + // bounded rescue policy, so it remains recoverable. + const rescueUntracked = join(dataDir, 'git-checkpoints', restored.rescueCheckpointId as string, 'untracked', 'huge.bin') + expect((await stat(rescueUntracked)).size).toBe(2_000_000) + }) + + it('fails closed before reset/clean when the rescue snapshot is partial', async () => { + await writeFile(join(repoRoot, 'huge.bin'), Buffer.alloc(6_000_000, 9)) + const checkpoint = await createGitCheckpoint({ + dataDir, + workspaceRoot: repoRoot, + threadId: 'thr_partial_rescue', + storage: { maxUntrackedFileBytes: 1_000_000 } + }) + if (!checkpoint.ok) throw new Error(checkpoint.message) + + const restored = await restoreGitCheckpoint({ + dataDir, + checkpointId: checkpoint.checkpointId, + allowPartialRestore: true + }) + expect(restored.ok).toBe(false) + if (restored.ok) throw new Error('expected incomplete rescue to refuse restore') + expect(restored.reason).toBe('partial') + expect((await stat(join(repoRoot, 'huge.bin'))).size).toBe(6_000_000) + }) + + it('prunes oldest checkpoints beyond the per-thread cap', async () => { + const ids: string[] = [] + for (let i = 0; i < 4; i += 1) { + const cp = await createGitCheckpoint({ + dataDir, + workspaceRoot: repoRoot, + threadId: 'thr_cap', + checkpointId: `gcp_${1000 + i}_fixed-${i}`, + storage: { maxPerThread: 2 } + }) + if (!cp.ok) throw new Error(cp.message) + ids.push(cp.checkpointId) + } + const root = join(dataDir, 'git-checkpoints') + // Only the two newest survive; the two oldest are pruned. + await expect(stat(join(root, ids[0]))).rejects.toBeTruthy() + await expect(stat(join(root, ids[1]))).rejects.toBeTruthy() + await expect(stat(join(root, ids[2]))).resolves.toBeTruthy() + await expect(stat(join(root, ids[3]))).resolves.toBeTruthy() + }) +}) diff --git a/src/main/services/git-checkpoint-service.ts b/src/main/services/git-checkpoint-service.ts index 562681eaa..823f598e7 100644 --- a/src/main/services/git-checkpoint-service.ts +++ b/src/main/services/git-checkpoint-service.ts @@ -17,8 +17,38 @@ type GitCheckpointMetadata = { currentBranch: string | null createdAt: string untrackedFiles: string[] + /** Untracked files deliberately NOT snapshotted (too large / over budget). */ + skippedUntracked?: string[] + /** + * Whether the snapshot captured every untracked file. `partial` means some + * untracked files were skipped (see `skippedUntracked`); restoring a partial + * checkpoint can destroy those never-captured files, so restore refuses a + * partial checkpoint unless the caller explicitly opts in. + */ + completeness?: 'complete' | 'partial' } +/** + * Snapshot policy that bounds checkpoint disk usage (issue #651). Untracked + * files are physically copied, so a workspace full of large untracked artifacts + * (AI models, node_modules, build output) could balloon the checkpoint store by + * gigabytes per message. These caps + the per-thread retention limit stop that. + */ +export type GitCheckpointStorageOptions = { + /** Override the checkpoints root (e.g. point it at another drive). */ + checkpointsRoot?: string + /** Skip snapshotting any single untracked file larger than this. Default 5 MiB. */ + maxUntrackedFileBytes?: number + /** Stop snapshotting untracked files once this cumulative size is reached. Default 50 MiB. */ + maxUntrackedTotalBytes?: number + /** Keep at most this many checkpoints per thread (newest first). Default 5. */ + maxPerThread?: number +} + +const DEFAULT_MAX_UNTRACKED_FILE_BYTES = 5 * 1_024 * 1_024 +const DEFAULT_MAX_UNTRACKED_TOTAL_BYTES = 50 * 1_024 * 1_024 +const DEFAULT_MAX_CHECKPOINTS_PER_THREAD = 5 + export type GitCheckpointCleanupResult = { scanned: number kept: number @@ -56,24 +86,38 @@ function restoreFailure(error: unknown): Extract { @@ -97,9 +141,9 @@ async function assertNoUnmerged(repositoryRoot: string): Promise { } } -async function readMetadata(dataDir: string, checkpointId: string): Promise { +async function readMetadata(root: string, checkpointId: string): Promise { try { - const raw = await readFile(metadataPath(dataDir, checkpointId), 'utf-8') + const raw = await readFile(metadataPath(root, checkpointId), 'utf-8') return JSON.parse(raw) as GitCheckpointMetadata } catch { return null @@ -133,13 +177,13 @@ async function writeHeadBundle(repositoryRoot: string, path: string): Promise { const head = metadata.head.trim() if (await commitExists(repositoryRoot, head)) return head - const bundlePath = checkpointHeadBundlePath(dataDir, metadata.checkpointId) + const bundlePath = checkpointHeadBundlePath(root, metadata.checkpointId) if (await fileExists(bundlePath)) { await runGit(repositoryRoot, ['bundle', 'unbundle', bundlePath], 30_000) if (await commitExists(repositoryRoot, head)) return head @@ -329,9 +373,9 @@ async function collectReferencedCheckpointIds(dataDir: string): Promise { +async function readCleanupState(root: string): Promise { try { - const raw = await readFile(checkpointCleanupStatePath(dataDir), 'utf-8') + const raw = await readFile(checkpointCleanupStatePath(root), 'utf-8') const parsed = JSON.parse(raw) as GitCheckpointCleanupState return typeof parsed === 'object' && parsed !== null ? parsed : {} } catch { @@ -339,10 +383,9 @@ async function readCleanupState(dataDir: string): Promise { - const root = checkpointRootDir(dataDir) +async function writeCleanupState(root: string, state: GitCheckpointCleanupState): Promise { await mkdir(root, { recursive: true }) - await writeFile(checkpointCleanupStatePath(dataDir), JSON.stringify(state, null, 2), 'utf-8') + await writeFile(checkpointCleanupStatePath(root), JSON.stringify(state, null, 2), 'utf-8') } function isCheckpointCleanupDue(lastRunAt: string | undefined, intervalDays: number, now: Date): boolean { @@ -362,12 +405,13 @@ const CHECKPOINT_CLEANUP_GRACE_MS = 10 * 60 * 1_000 export async function cleanupUnusedGitCheckpoints(params: { dataDir: string + checkpointsRoot?: string graceMs?: number now?: Date }): Promise { const graceMs = params.graceMs ?? CHECKPOINT_CLEANUP_GRACE_MS const nowMs = (params.now ?? new Date()).getTime() - const root = checkpointRootDir(params.dataDir) + const root = resolveCheckpointsRoot(params.dataDir, params.checkpointsRoot) const referenced = await collectReferencedCheckpointIds(params.dataDir) const result: GitCheckpointCleanupResult = { scanned: 0, @@ -422,19 +466,26 @@ export async function cleanupUnusedGitCheckpoints(params: { export async function cleanupUnusedGitCheckpointsIfDue(params: { dataDir: string + checkpointsRoot?: string intervalDays: number now?: Date graceMs?: number }): Promise { const now = params.now ?? new Date() - const state = await readCleanupState(params.dataDir) + const root = resolveCheckpointsRoot(params.dataDir, params.checkpointsRoot) + const state = await readCleanupState(root) const lastRunAt = typeof state.lastRunAt === 'string' ? state.lastRunAt : undefined if (!isCheckpointCleanupDue(lastRunAt, params.intervalDays, now)) { return { due: false, lastRunAt: lastRunAt ?? null } } - const result = await cleanupUnusedGitCheckpoints({ dataDir: params.dataDir, now, graceMs: params.graceMs }) + const result = await cleanupUnusedGitCheckpoints({ + dataDir: params.dataDir, + ...(params.checkpointsRoot ? { checkpointsRoot: params.checkpointsRoot } : {}), + now, + ...(params.graceMs !== undefined ? { graceMs: params.graceMs } : {}) + }) const nextLastRunAt = now.toISOString() - await writeCleanupState(params.dataDir, { lastRunAt: nextLastRunAt }) + await writeCleanupState(root, { lastRunAt: nextLastRunAt }) return { due: true, lastRunAt: nextLastRunAt, result } } @@ -443,11 +494,16 @@ export async function createGitCheckpoint(params: { workspaceRoot: string threadId: string checkpointId?: string + storage?: GitCheckpointStorageOptions }): Promise { const workspaceRoot = params.workspaceRoot.trim() if (!workspaceRoot) { return { ok: false, reason: 'no_workspace', message: 'No working directory selected.' } } + const root = resolveCheckpointsRoot(params.dataDir, params.storage?.checkpointsRoot) + const maxFileBytes = params.storage?.maxUntrackedFileBytes ?? DEFAULT_MAX_UNTRACKED_FILE_BYTES + const maxTotalBytes = params.storage?.maxUntrackedTotalBytes ?? DEFAULT_MAX_UNTRACKED_TOTAL_BYTES + const maxPerThread = params.storage?.maxPerThread ?? DEFAULT_MAX_CHECKPOINTS_PER_THREAD try { const repositoryRoot = await resolveRepositoryRoot(workspaceRoot) if (!repositoryRoot) { @@ -456,26 +512,47 @@ export async function createGitCheckpoint(params: { await assertNoUnmerged(repositoryRoot) const checkpointId = params.checkpointId?.trim() || `gcp_${Date.now()}_${randomUUID()}` - const dir = checkpointDir(params.dataDir, checkpointId) + const dir = checkpointDir(root, checkpointId) await rm(dir, { recursive: true, force: true }) await mkdir(join(dir, 'untracked'), { recursive: true }) const head = (await runGit(repositoryRoot, ['rev-parse', 'HEAD'])).stdout.trim() - await writeHeadBundle(repositoryRoot, checkpointHeadBundlePath(params.dataDir, checkpointId)) + await writeHeadBundle(repositoryRoot, checkpointHeadBundlePath(root, checkpointId)) const currentBranchRaw = (await runGit(repositoryRoot, ['branch', '--show-current'])).stdout.trim() const currentBranch = currentBranchRaw || null - const untrackedFiles = splitNul( + const candidateUntracked = splitNul( (await runGit(repositoryRoot, ['ls-files', '--others', '--exclude-standard', '-z'])).stdout ) await writePatch(repositoryRoot, ['diff', '--binary'], join(dir, 'unstaged.patch')) await writePatch(repositoryRoot, ['diff', '--cached', '--binary'], join(dir, 'staged.patch')) - for (const relativePath of untrackedFiles) { + // Bounded untracked snapshot (issue #651): copying every untracked file in + // full each turn is what ballooned the store by GBs. Skip files over the + // per-file cap and stop once the cumulative budget is hit; record what was + // skipped so the model/user know the snapshot is partial. + const untrackedFiles: string[] = [] + const skippedUntracked: string[] = [] + let copiedBytes = 0 + for (const relativePath of candidateUntracked) { const from = join(repositoryRoot, relativePath) + let size = 0 + try { + const info = await stat(from) + if (info.isDirectory()) continue + size = info.size + } catch { + continue + } + if (size > maxFileBytes || copiedBytes + size > maxTotalBytes) { + skippedUntracked.push(relativePath) + continue + } const to = join(dir, 'untracked', relativePath) await mkdir(dirname(to), { recursive: true }) await cp(from, to, { recursive: true, force: true, errorOnExist: false }) + copiedBytes += size + untrackedFiles.push(relativePath) } const metadata: GitCheckpointMetadata = { @@ -485,9 +562,13 @@ export async function createGitCheckpoint(params: { head, currentBranch, createdAt: new Date().toISOString(), - untrackedFiles + untrackedFiles, + ...(skippedUntracked.length ? { skippedUntracked } : {}), + completeness: skippedUntracked.length ? 'partial' : 'complete' } await writeFile(join(dir, 'metadata.json'), JSON.stringify(metadata, null, 2), 'utf-8') + // Bound per-thread retention so an active thread cannot grow unboundedly. + await pruneThreadCheckpoints(root, params.threadId, maxPerThread, checkpointId).catch(() => undefined) return { ok: true, checkpointId, repositoryRoot, head, currentBranch } } catch (error) { const failure = checkpointFailure(error) @@ -498,9 +579,69 @@ export async function createGitCheckpoint(params: { } } +/** + * Keep at most `max` checkpoints for a thread (issue #651, per-thread cap). + * Oldest checkpoints (by createdAt, falling back to the `gcp__` name) are + * removed first; `keepId` (the just-created checkpoint) is always retained. + */ +export async function pruneThreadCheckpoints( + root: string, + threadId: string, + max: number, + keepId?: string +): Promise<{ deleted: string[] }> { + if (max <= 0) return { deleted: [] } + let entries: Dirent[] + try { + entries = await readdir(root, { withFileTypes: true }) + } catch { + return { deleted: [] } + } + const owned: Array<{ id: string; order: number }> = [] + for (const entry of entries) { + if (!entry.isDirectory()) continue + const metadata = await readMetadata(root, entry.name) + if (!metadata || metadata.threadId !== threadId) continue + const createdMs = Date.parse(metadata.createdAt) + const order = Number.isFinite(createdMs) ? createdMs : checkpointNameTimestamp(entry.name) + owned.push({ id: entry.name, order }) + } + // Newest first; keep the first `max`, delete the rest (never the keepId). + owned.sort((a, b) => b.order - a.order) + const deleted: string[] = [] + for (let i = 0; i < owned.length; i += 1) { + const { id } = owned[i] + if (i < max || id === keepId) continue + try { + await rm(checkpointDir(root, id), { recursive: true, force: true }) + deleted.push(id) + } catch { + // best-effort + } + } + return { deleted } +} + +/** Extract the `gcp__` creation epoch for ordering fallback. */ +function checkpointNameTimestamp(name: string): number { + const match = name.match(/^gcp_(\d+)_/) + return match ? Number(match[1]) : 0 +} + export async function restoreGitCheckpoint(params: { dataDir: string checkpointId: string + storage?: GitCheckpointStorageOptions + /** + * Opt-in to restoring a PARTIAL checkpoint (one whose snapshot skipped some + * untracked files because they were over the size budget). A partial restore + * runs `git clean -fd`, which deletes those never-captured files; without this + * flag the restore is refused so the user does not silently lose data. When + * enabled, a complete rescue checkpoint is taken first so the cleaned files + * remain recoverable. The restore still fails closed when the configured + * checkpoint budget cannot capture that rescue. + */ + allowPartialRestore?: boolean /** * Optional runtime bridge used to verify that no thread is mid-turn before * running the destructive `git reset --hard` / `git clean -fd`. When omitted @@ -511,14 +652,32 @@ export async function restoreGitCheckpoint(params: { runtimeRequest?: (path: string, init: { method?: string; body?: string }) => Promise<{ ok: boolean; status: number; body: string }> }): Promise { const checkpointId = params.checkpointId.trim() - const metadata = await readMetadata(params.dataDir, checkpointId) + const root = resolveCheckpointsRoot(params.dataDir, params.storage?.checkpointsRoot) + const metadata = await readMetadata(root, checkpointId) if (!metadata) { return { ok: false, reason: 'not_found', message: `Git checkpoint not found: ${checkpointId}` } } + + // Partial-checkpoint data-loss guard (P0-01). If the snapshot skipped any + // untracked file, the upcoming `git clean -fd` would delete those files with + // no snapshot to restore them. Refuse unless the caller explicitly opts in. + const skippedUntracked = metadata.skippedUntracked ?? [] + const isPartial = metadata.completeness === 'partial' || skippedUntracked.length > 0 + if (isPartial && !params.allowPartialRestore) { + return { + ok: false, + reason: 'partial', + message: + `This checkpoint is partial: ${skippedUntracked.length} untracked file(s) were too large to snapshot and are NOT stored in it. ` + + 'Restoring would permanently delete them. Re-run with allowPartialRestore to proceed (a full rescue checkpoint will be taken first).', + skippedUntracked + } + } + try { const repositoryRoot = metadata.repositoryRoot await assertNoUnmerged(repositoryRoot) - const targetRef = await resolveCheckpointTarget(repositoryRoot, params.dataDir, metadata) + const targetRef = await resolveCheckpointTarget(repositoryRoot, root, metadata) // Busy guard: a checkpoint restore runs `git reset --hard` + `git clean // -fd`, which would destroy files the agent is actively editing. Before @@ -560,12 +719,35 @@ export async function restoreGitCheckpoint(params: { } } + // The rescue checkpoint is the safety net for `reset --hard` + `clean -fd`. + // Never bypass the storage budget here: an unbounded rescue reintroduces the + // disk-exhaustion failure this service is meant to prevent. Instead require a + // COMPLETE rescue and fail closed before the first destructive git command. const rescue = await createGitCheckpoint({ dataDir: params.dataDir, + storage: params.storage, workspaceRoot: repositoryRoot, threadId: `${metadata.threadId}:rollback-rescue` }) - const rescueCheckpointId = rescue.ok ? rescue.checkpointId : null + if (!rescue.ok) { + return { + ok: false, + reason: rescue.reason, + message: `Cannot safely restore checkpoint because the rescue snapshot failed: ${rescue.message}` + } + } + const rescueMetadata = await readMetadata(root, rescue.checkpointId) + if (!rescueMetadata || rescueMetadata.completeness !== 'complete') { + return { + ok: false, + reason: 'partial', + message: + 'Cannot safely restore checkpoint because the current workspace does not fit the configured rescue snapshot limits. ' + + 'Increase the checkpoint limits or move/remove the oversized untracked files, then retry.', + skippedUntracked: rescueMetadata?.skippedUntracked ?? [] + } + } + const rescueCheckpointId = rescue.checkpointId await runGit(repositoryRoot, ['reset', '--hard'], 30_000) await runGit(repositoryRoot, ['clean', '-fd'], 30_000) @@ -577,7 +759,7 @@ export async function restoreGitCheckpoint(params: { await runGit(repositoryRoot, ['reset', '--hard', targetRef], 30_000) await runGit(repositoryRoot, ['clean', '-fd'], 30_000) - const dir = checkpointDir(params.dataDir, checkpointId) + const dir = checkpointDir(root, checkpointId) await applyPatchIfPresent(repositoryRoot, join(dir, 'staged.patch'), true) await applyPatchIfPresent(repositoryRoot, join(dir, 'unstaged.patch'), false) diff --git a/src/main/settings-store.test.ts b/src/main/settings-store.test.ts index bda167787..e497da9c5 100644 --- a/src/main/settings-store.test.ts +++ b/src/main/settings-store.test.ts @@ -4,6 +4,7 @@ import { join } from 'node:path' import { describe, expect, it } from 'vitest' import { DEFAULT_APPROVAL_POLICY, + DEFAULT_CHECKPOINT_CLEANUP_ENABLED, DEFAULT_CHECKPOINT_CLEANUP_INTERVAL_DAYS, defaultKunRuntimeSettings, defaultModelProviderSettings @@ -21,8 +22,8 @@ describe('JsonSettingsStore', () => { expect(loaded.guiUpdate.channel).toBe(DEFAULT_GUI_UPDATE_CHANNEL) expect(loaded.agents.kun.approvalPolicy).toBe(DEFAULT_APPROVAL_POLICY) expect(loaded.checkpointCleanup.intervalDays).toBe(DEFAULT_CHECKPOINT_CLEANUP_INTERVAL_DAYS) - // Checkpoint cleanup deletes data, so it must be opt-in (disabled by default). - expect(loaded.checkpointCleanup.enabled).toBe(false) + // Checkpoint cleanup is enabled by default to keep stale checkpoints from accumulating. + expect(loaded.checkpointCleanup.enabled).toBe(DEFAULT_CHECKPOINT_CLEANUP_ENABLED) expect(loaded.appBehavior).toEqual({ openAtLogin: false, startMinimized: false, diff --git a/src/main/settings-store.ts b/src/main/settings-store.ts index 0721a8428..9cfc26e2c 100644 --- a/src/main/settings-store.ts +++ b/src/main/settings-store.ts @@ -9,6 +9,7 @@ import { DEFAULT_CHECKPOINT_CLEANUP_ENABLED, DEFAULT_CHECKPOINT_CLEANUP_INTERVAL_DAYS, DEFAULT_CURSOR_SPOTLIGHT_COLOR, + DEFAULT_GIT_BRANCH_PREFIX, DEFAULT_LOG_RETENTION_DAYS, DEFAULT_WRITE_WORKSPACE_ROOT, defaultClawSettings, @@ -31,6 +32,7 @@ import { DEFAULT_UI_FONT_SCALE, normalizeAppBehaviorSettings, normalizeCheckpointCleanupSettings, + normalizeGitBranchPrefix, normalizeKeyboardShortcuts, migrateLegacyAppSettings, normalizeAppSettings, @@ -241,6 +243,7 @@ const defaultSettings = (): AppSettingsV1 => ({ enabled: DEFAULT_CHECKPOINT_CLEANUP_ENABLED, intervalDays: DEFAULT_CHECKPOINT_CLEANUP_INTERVAL_DAYS }, + gitBranchPrefix: DEFAULT_GIT_BRANCH_PREFIX, notifications: { turnComplete: true }, @@ -273,6 +276,7 @@ function buildMergedSettings(parsed: Partial): AppSettingsV1 { ...defaults.checkpointCleanup, ...migrated.checkpointCleanup }), + gitBranchPrefix: normalizeGitBranchPrefix(migrated.gitBranchPrefix), notifications: { ...defaults.notifications, ...migrated.notifications }, appBehavior: mergeAppBehaviorSettings(defaults.appBehavior, migrated.appBehavior), keyboardShortcuts: normalizeKeyboardShortcuts(migrated.keyboardShortcuts), diff --git a/src/renderer/src/agent/kun-contract.ts b/src/renderer/src/agent/kun-contract.ts index cb73c6fbc..4dafc3a54 100644 --- a/src/renderer/src/agent/kun-contract.ts +++ b/src/renderer/src/agent/kun-contract.ts @@ -53,11 +53,15 @@ export type CoreThreadJson = CoreThreadSummaryJson & { export type CoreAttachmentMetadataJson = { id: string name: string + kind?: 'image' | 'document' mimeType: string byteSize: number hash: string width?: number height?: number + documentText?: string + pageCount?: number + truncated?: boolean localFilePath?: string textFallback?: CoreAttachmentTextFallbackJson threadIds?: string[] @@ -223,6 +227,9 @@ export type CoreRuntimeCapabilityManifestJson = { maxImageBytes: number maxImageDimension: number allowedMimeTypes: string[] + allowedDocumentMimeTypes?: string[] + maxDocumentBytes?: number + maxDocumentTextChars?: number textFallbackMaxBase64Bytes?: number textFallbackMaxImageDimension?: number textFallbackPreferredMimeType?: string @@ -274,6 +281,7 @@ export type CoreRuntimeInfoJson = { export type CoreRuntimeToolDiagnosticsJson = { providers?: Array> mcpServers?: Array> + mcpOAuth?: CoreMcpOAuthDiagnosticJson[] mcpSearch?: { enabled?: boolean mode?: 'direct' | 'search' | 'auto' @@ -305,6 +313,38 @@ export type CoreRuntimeToolDiagnosticsJson = { } } +export type CoreMcpOAuthDiagnosticJson = { + serverId: string + enabled: boolean + configured: boolean + transport: string + url?: string + status: 'disabled' | 'empty' | 'partial' | 'authorized' | 'expired' | 'error' + hasClientInformation: boolean + hasTokens: boolean + hasRefreshToken: boolean + hasCodeVerifier: boolean + hasDiscoveryState: boolean + grantedScopes?: string[] + expiresAt?: string + lastError?: string + lastErrorAt?: string +} + +export type CoreMcpOAuthDiagnosticsResponseJson = { + servers: CoreMcpOAuthDiagnosticJson[] +} + +export type CoreMcpOAuthClearResponseJson = { + cleared: string[] +} + +export type CoreMcpOAuthAuthorizeResponseJson = { + serverId: string + status: CoreMcpOAuthDiagnosticJson['status'] + authorized: boolean +} + export type CoreRuntimeSkillJson = { id: string name: string diff --git a/src/renderer/src/agent/kun-runtime.test.ts b/src/renderer/src/agent/kun-runtime.test.ts index e07fbf946..6815ce277 100644 --- a/src/renderer/src/agent/kun-runtime.test.ts +++ b/src/renderer/src/agent/kun-runtime.test.ts @@ -96,6 +96,20 @@ describe('KunRuntimeProvider', () => { ) }) + it('starts MCP OAuth authorization through the authenticated runtime bridge', async () => { + const runtimeRequest = vi.fn(async () => ({ + ok: true, + status: 200, + body: JSON.stringify({ serverId: 'google_drive', status: 'authorized', authorized: true }) + })) + installDsGui({ runtimeRequest }) + + const result = await new KunRuntimeProvider().authorizeMcpOAuthCredentials('google_drive') + + expect(runtimeRequest).toHaveBeenCalledWith('/v1/mcp/oauth/google_drive', 'POST') + expect(result).toEqual({ serverId: 'google_drive', status: 'authorized', authorized: true }) + }) + it('maps Kun thread items into chat blocks', async () => { installDsGui({ runtimeRequest: vi.fn(async () => ({ @@ -648,6 +662,28 @@ describe('KunRuntimeProvider', () => { threadId: 'thr_1' }) ) + await expect(provider.uploadAttachment({ + name: 'spec.pdf', + mimeType: 'application/pdf', + dataBase64: 'JVBERi0=', + documentText: 'PDF body', + pageCount: 2, + localFilePath: '/tmp/picked/spec.pdf', + workspace: '/tmp/ws' + })).resolves.toMatchObject({ id: 'att_1' }) + expect(runtimeRequest).toHaveBeenCalledWith( + '/v1/attachments', + 'POST', + JSON.stringify({ + name: 'spec.pdf', + mimeType: 'application/pdf', + dataBase64: 'JVBERi0=', + documentText: 'PDF body', + pageCount: 2, + localFilePath: '/tmp/picked/spec.pdf', + workspace: '/tmp/ws' + }) + ) }) it('lists, disables, and deletes memory records through Kun endpoints', async () => { diff --git a/src/renderer/src/agent/kun-runtime.ts b/src/renderer/src/agent/kun-runtime.ts index f521a70b3..49b5803af 100644 --- a/src/renderer/src/agent/kun-runtime.ts +++ b/src/renderer/src/agent/kun-runtime.ts @@ -14,6 +14,7 @@ import { KUN_ATTACHMENTS_PATH, KUN_MEMORY_DIAGNOSTICS_PATH, KUN_MEMORY_PATH, + KUN_MCP_OAUTH_PATH, KUN_RUNTIME_INFO_PATH, KUN_RUNTIME_TOOLS_PATH, KUN_SKILLS_PATH, @@ -32,6 +33,7 @@ import { kunAttachmentContentPath, kunUserInputPath, kunMemoryRecordPath, + kunMcpOAuthServerPath, kunSessionResumePath, normalizeThreadMode, type KunThreadMode @@ -46,6 +48,10 @@ import type { CoreMemoryDiagnosticsJson, CoreMemoryListResponseJson, CoreMemoryRecordJson, + CoreMcpOAuthClearResponseJson, + CoreMcpOAuthAuthorizeResponseJson, + CoreMcpOAuthDiagnosticJson, + CoreMcpOAuthDiagnosticsResponseJson, CoreResumeSessionResponseJson, CoreRuntimeInfoJson, CoreRuntimeEventJson, @@ -640,6 +646,42 @@ export class KunRuntimeProvider implements AgentProvider { ) } + async getMcpOAuthDiagnostics(): Promise { + const response = await rendererRuntimeClient.runtimeRequest(KUN_MCP_OAUTH_PATH, 'GET') + if (!response.ok) { + throw runtimeErrorToError(readRuntimeError(response.body, 'failed to load MCP OAuth diagnostics')) + } + return readRuntimeJson( + response.body, + 'runtime returned an invalid MCP OAuth diagnostics response' + ).servers + } + + async clearMcpOAuthCredentials(serverId?: string): Promise { + const response = await rendererRuntimeClient.runtimeRequest( + serverId ? kunMcpOAuthServerPath(serverId) : KUN_MCP_OAUTH_PATH, + 'DELETE' + ) + if (!response.ok) { + throw runtimeErrorToError(readRuntimeError(response.body, 'failed to clear MCP OAuth credentials')) + } + return readRuntimeJson( + response.body, + 'runtime returned an invalid MCP OAuth reset response' + ).cleared + } + + async authorizeMcpOAuthCredentials(serverId: string): Promise { + const response = await rendererRuntimeClient.runtimeRequest(kunMcpOAuthServerPath(serverId), 'POST') + if (!response.ok) { + throw runtimeErrorToError(readRuntimeError(response.body, 'failed to authorize MCP OAuth connector')) + } + return readRuntimeJson( + response.body, + 'runtime returned an invalid MCP OAuth authorize response' + ) + } + async listSkills(): Promise { const response = await rendererRuntimeClient.runtimeRequest(KUN_SKILLS_PATH, 'GET') if (!response.ok) { @@ -655,6 +697,8 @@ export class KunRuntimeProvider implements AgentProvider { name: string mimeType?: string dataBase64: string + documentText?: string + pageCount?: number localFilePath?: string textFallback?: CoreAttachmentTextFallbackJson threadId?: string diff --git a/src/renderer/src/agent/types.ts b/src/renderer/src/agent/types.ts index bd65b4670..f43b10bf2 100644 --- a/src/renderer/src/agent/types.ts +++ b/src/renderer/src/agent/types.ts @@ -4,6 +4,7 @@ import type { CoreAttachmentTextFallbackJson, CoreMemoryDiagnosticsJson, CoreMemoryRecordJson, + CoreMcpOAuthDiagnosticJson, CoreRuntimeInfoJson, CoreRuntimeSkillJson, CoreRuntimeToolDiagnosticsJson @@ -500,11 +501,16 @@ export interface AgentProvider { ): Promise<{ turnId: string; threadId: string; userMessageItemId?: string; reviewItemId?: string }> getRuntimeInfo?(): Promise getToolDiagnostics?(): Promise + getMcpOAuthDiagnostics?(): Promise + clearMcpOAuthCredentials?(serverId?: string): Promise + authorizeMcpOAuthCredentials?(serverId: string): Promise listSkills?(): Promise uploadAttachment?(input: { name: string mimeType?: string dataBase64: string + documentText?: string + pageCount?: number localFilePath?: string textFallback?: CoreAttachmentTextFallbackJson threadId?: string diff --git a/src/renderer/src/components/PluginMarketplaceView.test.ts b/src/renderer/src/components/PluginMarketplaceView.test.ts index fc1dd4f68..ebd3792a2 100644 --- a/src/renderer/src/components/PluginMarketplaceView.test.ts +++ b/src/renderer/src/components/PluginMarketplaceView.test.ts @@ -4,6 +4,8 @@ import { buildMcpConfig, buildRemoteMcpConfig, customMcpConfigFragment, + auditMcpConfigSupplyChain, + auditMarketplaceInstall, isAllowedDocsUrl, isHttpsUrl, mcpConfigHasServer, @@ -84,7 +86,7 @@ describe('PluginMarketplaceView MCP config helpers', () => { const merged = mergeMcpJsonConfig( existing, - buildMcpConfig('playwright', 'npx', ['-y', '@playwright/mcp@latest']) + buildMcpConfig('playwright', 'npx', ['-y', '@playwright/mcp@0.0.77']) ) const parsed = JSON.parse(merged.text) as Record @@ -95,14 +97,14 @@ describe('PluginMarketplaceView MCP config helpers', () => { enabled: true, transport: 'stdio', command: 'npx', - args: ['-y', '@playwright/mcp@latest'], + args: ['-y', '@playwright/mcp@0.0.77'], trustScope: 'user' }) expect(mcpConfigHasServer(merged.text, 'playwright')).toBe(true) }) it('detects duplicate MCP servers instead of appending old-style snippets', () => { - const fragment = buildMcpConfig('context7', 'npx', ['-y', '@upstash/context7-mcp@latest']) + const fragment = buildMcpConfig('context7', 'npx', ['-y', '@upstash/context7-mcp@3.2.2']) const first = mergeMcpJsonConfig('', fragment) const second = mergeMcpJsonConfig(first.text, fragment) @@ -111,6 +113,36 @@ describe('PluginMarketplaceView MCP config helpers', () => { expect(JSON.parse(second.text).servers.context7).toMatchObject({ command: 'npx' }) }) + it.each([ + ['npx', '@upstash/context7-mcp@latest'], + ['npx.cmd', '@upstash/context7-mcp@3.2'], + ['C:\\Program Files\\nodejs\\npx.cmd', '@upstash/context7-mcp@^3.2.2'] + ])('rejects non-exact package versions for %s', (command, packageSpec) => { + const audit = auditMcpConfigSupplyChain(buildMcpConfig('context7', command, ['-y', packageSpec])) + expect(audit.ok).toBe(false) + expect(audit.errors.join('\n')).toMatch(/exact npm package version/) + }) + + it.each(['npx', 'npx.cmd', 'C:\\Program Files\\nodejs\\npx.cmd'])('allows exact package versions for %s', (command) => { + const audit = auditMcpConfigSupplyChain(buildMcpConfig('context7', command, ['-y', '@upstash/context7-mcp@3.2.2'])) + expect(audit.ok).toBe(true) + expect(audit.permissions).toEqual(expect.arrayContaining(['command', 'file', 'network'])) + }) + + it('audits a marketplace MCP install before writing config', () => { + const audit = auditMarketplaceInstall({ + id: 'context7', + kind: 'mcp', + title: 'Context7', + description: 'Use Context7', + group: 'recommended', + mcpConfig: () => buildMcpConfig('context7', 'npx', ['-y', '@upstash/context7-mcp@3.2.2']), + supplyChain: { source: 'mcp', permissions: ['command', 'network'] } + }, '/workspace') + expect(audit.ok).toBe(true) + expect(audit.permissions).toEqual(expect.arrayContaining(['command', 'file', 'network'])) + }) + it('accepts custom JSON as either a single server or a Kun config fragment', () => { expect(customMcpConfigFragment( 'docs', diff --git a/src/renderer/src/components/PluginMarketplaceView.tsx b/src/renderer/src/components/PluginMarketplaceView.tsx index df06a6b91..fb1e92ec2 100644 --- a/src/renderer/src/components/PluginMarketplaceView.tsx +++ b/src/renderer/src/components/PluginMarketplaceView.tsx @@ -62,11 +62,27 @@ type MarketplaceItem = { serverIds?: string[] mcpConfig?: (workspaceRoot: string) => JsonRecord oauth?: OAuthConnectorInfo + supplyChain?: SupplyChainInfo skillInstructions?: string } type JsonRecord = Record +type SupplyChainPermission = 'file' | 'command' | 'network' | 'secret' + +type SupplyChainInfo = { + source: 'mcp' | 'remote-mcp' | 'skill' + permissions: SupplyChainPermission[] + packageName?: string + version?: string +} + +type MarketplaceInstallAudit = { + ok: boolean + permissions: SupplyChainPermission[] + errors: string[] +} + type OAuthConnectorInfo = { docsUrl: string permissionKeys: string[] @@ -249,6 +265,90 @@ function buildRemoteMcpServer(url: string): JsonRecord { } } +const PERMISSION_LABELS: Record = { + file: 'File', + command: 'Command', + network: 'Network', + secret: 'Secret' +} + +function uniquePermissions(permissions: readonly SupplyChainPermission[]): SupplyChainPermission[] { + return [...new Set(permissions)] +} + +export function auditMcpConfigSupplyChain(config: JsonRecord): MarketplaceInstallAudit { + const permissions: SupplyChainPermission[] = ['network'] + const errors: string[] = [] + const servers = isJsonRecord(config.servers) ? config.servers : {} + for (const [id, server] of Object.entries(servers)) { + if (!isJsonRecord(server)) { + errors.push(`MCP server "${id}" must be an object.`) + continue + } + if (typeof server.command === 'string' && server.command.trim()) { + permissions.push('command', 'file') + const command = executableBasename(server.command) + const args = Array.isArray(server.args) ? server.args.filter((arg): arg is string => typeof arg === 'string') : [] + if (command === 'npx') { + const spec = npxPackageSpec(args) + if (!spec) { + errors.push(`MCP server "${id}" uses npx without an exact package version.`) + } else if (!isExactNpmPackageSpec(spec)) { + errors.push(`MCP server "${id}" must pin an exact npm package version (got "${spec}").`) + } + } + } + if (isJsonRecord(server.env) && Object.keys(server.env).length > 0) { + permissions.push('secret') + } + if (typeof server.url === 'string') { + permissions.push('network') + if (!isHttpsUrl(server.url)) errors.push(`MCP server "${id}" URL must use HTTPS.`) + } + } + return { ok: errors.length === 0, permissions: uniquePermissions(permissions), errors } +} + +export function auditMarketplaceInstall(item: MarketplaceItem, workspaceRoot: string): MarketplaceInstallAudit { + const configuredPermissions = item.supplyChain?.permissions ?? [] + const config = item.mcpConfig?.(workspaceRoot) + if (!config) { + return { + ok: false, + permissions: uniquePermissions(configuredPermissions), + errors: ['MCP config is missing.'] + } + } + const audit = auditMcpConfigSupplyChain(config) + return { + ok: audit.ok, + permissions: uniquePermissions([...configuredPermissions, ...audit.permissions]), + errors: audit.errors + } +} + +function executableBasename(command: string): string { + const normalized = command.trim().replace(/^['"]|['"]$/g, '').replace(/\\/g, '/') + const basename = normalized.slice(normalized.lastIndexOf('/') + 1).toLowerCase() + return basename.replace(/\.(?:cmd|exe)$/i, '') +} + +function npxPackageSpec(args: readonly string[]): string | null { + for (const arg of args) { + if (!arg || arg.startsWith('-')) continue + return arg + } + return null +} + +function isExactNpmPackageSpec(spec: string): boolean { + if (spec.endsWith('@latest') || spec.includes('*') || spec.includes('^') || spec.includes('~')) return false + const at = spec.lastIndexOf('@') + if (at <= 0) return false + const version = spec.slice(at + 1) + return /^\d+\.\d+\.\d+(?:-[0-9A-Za-z.-]+)?(?:\+[0-9A-Za-z.-]+)?$/.test(version) +} + export function buildMcpConfig( id: string, command: string, @@ -594,8 +694,9 @@ const RECOMMENDED_ITEMS: MarketplaceItem[] = [ buildMcpConfig( 'playwright', 'npx', - ['-y', '@playwright/mcp@latest'] - ) + ['-y', '@playwright/mcp@0.0.77'] + ), + supplyChain: { source: 'mcp', packageName: '@playwright/mcp', version: '0.0.77', permissions: ['command', 'network', 'file'] } }, { id: 'github', @@ -607,8 +708,9 @@ const RECOMMENDED_ITEMS: MarketplaceItem[] = [ buildMcpConfig( 'github', 'npx', - ['-y', '@modelcontextprotocol/server-github'] - ) + ['-y', '@modelcontextprotocol/server-github@2025.4.8'] + ), + supplyChain: { source: 'mcp', packageName: '@modelcontextprotocol/server-github', version: '2025.4.8', permissions: ['command', 'network', 'secret'] } }, { id: 'vercel', @@ -635,6 +737,7 @@ const RECOMMENDED_ITEMS: MarketplaceItem[] = [ ], noteKey: 'pluginOAuthVercelNote' }, + supplyChain: { source: 'remote-mcp', permissions: ['network', 'secret'] }, mcpConfig: () => buildRemoteMcpConfig({ vercel: 'https://mcp.vercel.com' @@ -669,6 +772,7 @@ const RECOMMENDED_ITEMS: MarketplaceItem[] = [ ], noteKey: 'pluginOAuthGoogleNote' }, + supplyChain: { source: 'remote-mcp', permissions: ['network', 'secret', 'file'] }, mcpConfig: () => buildRemoteMcpConfig(GOOGLE_WORKSPACE_MCP_SERVERS) }, @@ -682,8 +786,9 @@ const RECOMMENDED_ITEMS: MarketplaceItem[] = [ buildMcpConfig( 'context7', 'npx', - ['-y', '@upstash/context7-mcp@latest'] - ) + ['-y', '@upstash/context7-mcp@3.2.2'] + ), + supplyChain: { source: 'mcp', packageName: '@upstash/context7-mcp', version: '3.2.2', permissions: ['command', 'network'] } }, { id: 'sequential-thinking', @@ -695,8 +800,9 @@ const RECOMMENDED_ITEMS: MarketplaceItem[] = [ buildMcpConfig( 'sequential-thinking', 'npx', - ['-y', '@modelcontextprotocol/server-sequential-thinking'] - ) + ['-y', '@modelcontextprotocol/server-sequential-thinking@2025.12.18'] + ), + supplyChain: { source: 'mcp', packageName: '@modelcontextprotocol/server-sequential-thinking', version: '2025.12.18', permissions: ['command'] } }, { id: 'memory', @@ -708,8 +814,9 @@ const RECOMMENDED_ITEMS: MarketplaceItem[] = [ buildMcpConfig( 'memory', 'npx', - ['-y', '@modelcontextprotocol/server-memory'] - ) + ['-y', '@modelcontextprotocol/server-memory@2026.1.26'] + ), + supplyChain: { source: 'mcp', packageName: '@modelcontextprotocol/server-memory', version: '2026.1.26', permissions: ['command', 'file'] } }, { id: 'brave-search', @@ -721,9 +828,10 @@ const RECOMMENDED_ITEMS: MarketplaceItem[] = [ buildMcpConfig( 'brave-search', 'npx', - ['-y', '@modelcontextprotocol/server-brave-search'], + ['-y', '@modelcontextprotocol/server-brave-search@0.6.2'], { env: { BRAVE_API_KEY: '' } } - ) + ), + supplyChain: { source: 'mcp', packageName: '@modelcontextprotocol/server-brave-search', version: '0.6.2', permissions: ['command', 'network', 'secret'] } }, { id: 'code-review', @@ -731,6 +839,7 @@ const RECOMMENDED_ITEMS: MarketplaceItem[] = [ titleKey: 'pluginSkillReviewTitle', descriptionKey: 'pluginSkillReviewDesc', group: 'recommended', + supplyChain: { source: 'skill', permissions: ['file'] }, skillInstructions: 'Use this skill when reviewing a code change. Prioritize correctness, regressions, security, performance, and missing tests. Lead with concrete findings and file references.' }, @@ -740,6 +849,7 @@ const RECOMMENDED_ITEMS: MarketplaceItem[] = [ titleKey: 'pluginSkillFrontendTitle', descriptionKey: 'pluginSkillFrontendDesc', group: 'recommended', + supplyChain: { source: 'skill', permissions: ['file'] }, skillInstructions: 'Use this skill when improving UI. Preserve the product style, check responsive states, avoid generic layouts, and verify the result visually before handing it back.' }, @@ -749,6 +859,7 @@ const RECOMMENDED_ITEMS: MarketplaceItem[] = [ titleKey: 'pluginSkillBugTitle', descriptionKey: 'pluginSkillBugDesc', group: 'recommended', + supplyChain: { source: 'skill', permissions: ['file', 'command'] }, skillInstructions: 'Use this skill when investigating bugs. Reproduce or narrow the symptom, trace the data flow, identify the smallest fix, and add focused verification where possible.' }, @@ -758,6 +869,7 @@ const RECOMMENDED_ITEMS: MarketplaceItem[] = [ titleKey: 'pluginSkillReleaseTitle', descriptionKey: 'pluginSkillReleaseDesc', group: 'recommended', + supplyChain: { source: 'skill', permissions: ['file'] }, skillInstructions: 'Use this skill when preparing release notes. Group user-facing changes by outcome, call out migrations or risks, and keep wording concise and scannable.' } @@ -1053,6 +1165,8 @@ export function PluginMarketplaceView({ leftSidebarCollapsed, onToggleLeftSideba const installMcpItem = async (item: MarketplaceItem): Promise => { if (!item.mcpConfig) return + const audit = auditMarketplaceInstall(item, workspaceRoot) + if (!audit.ok) throw new Error(audit.errors.join('\n')) await appendMcpConfig(item.id, item.mcpConfig(workspaceRoot)) } @@ -1130,7 +1244,10 @@ export function PluginMarketplaceView({ leftSidebarCollapsed, onToggleLeftSideba .map((arg) => arg.trim()) .filter(Boolean) ) - await appendMcpConfig(id, customMcpConfigFragment(id, customConfig, fallback)) + const fragment = customMcpConfigFragment(id, customConfig, fallback) + const audit = auditMcpConfigSupplyChain(fragment) + if (!audit.ok) throw new Error(audit.errors.join('\n')) + await appendMcpConfig(id, fragment) } else { if (!selectedSkillRoot?.path) { setNotice({ tone: 'error', message: t('pluginSkillRootMissing') }) @@ -1898,6 +2015,11 @@ function PluginSection({ {t('pluginOAuthBadge')} ) : null} + {item.supplyChain?.permissions.length ? ( + + {item.supplyChain.permissions.map((permission) => PERMISSION_LABELS[permission]).join(' / ')} + + ) : null}

{itemDescription(item, t)} diff --git a/src/renderer/src/components/SessionHeader.tsx b/src/renderer/src/components/SessionHeader.tsx index ed2454135..4f0cf4a3b 100644 --- a/src/renderer/src/components/SessionHeader.tsx +++ b/src/renderer/src/components/SessionHeader.tsx @@ -10,7 +10,7 @@ import { formatCacheMissReason, formatCost, formatPercent, - primaryCacheHitRate, + cumulativeCacheHitRate, useThreadUsage } from '../hooks/use-thread-usage' @@ -214,7 +214,7 @@ export function SessionHeader({ compact = false, className = '' }: Props): React } )} > - {t('sessionUsageCache', { cache: formatPercent(primaryCacheHitRate(threadUsage)) })} + {t('sessionUsageCache', { cache: formatPercent(cumulativeCacheHitRate(threadUsage)) })} ) : null} diff --git a/src/renderer/src/components/Workbench.tsx b/src/renderer/src/components/Workbench.tsx index d733ab94a..ff63a9dab 100644 --- a/src/renderer/src/components/Workbench.tsx +++ b/src/renderer/src/components/Workbench.tsx @@ -60,11 +60,17 @@ import { composeWritePrompt } from '../write/quoted-selection' import { resolveWriteAgentPreset } from '../write/agent-presets' import { useWriteWorkspaceStore } from '../write/write-workspace-store' import { isWriteThreadId } from '../write/write-thread-registry' -import { buildSddDraftId, createSddDraft, forgetRememberedSddDraft, useSddDraftStore } from '../sdd/sdd-draft-store' +import { + buildSddDraftId, + createSddDraft, + forgetRememberedSddDraft, + resolveSddRequirementWorkspace, + useSddDraftStore +} from '../sdd/sdd-draft-store' import type { SddDraft, SddDraftSaveStatus } from '../sdd/sdd-draft-store' import { listSddDraftHistory, titleFromSddDraftContent } from '../sdd/sdd-draft-history' import { saveActiveSddDraftToDisk } from '../sdd/sdd-draft-actions' -import { restoreRememberedSddDraft, restoreSddDraft } from '../sdd/sdd-draft-restore' +import { restoreSddDraft } from '../sdd/sdd-draft-restore' import { composeSddAssistantPrompt } from '../sdd/sdd-assistant-prompt' import { frameworkById } from '../sdd/pm-skill-frameworks' import { collectSddDraftImages, withAttachmentIds, type SddDraftImageReference } from '../sdd/sdd-draft-images' @@ -90,7 +96,7 @@ import { CODE_PANEL_PREFERRED, useWorkbenchLayout } from './workbench-layout' import { useWorkbenchPlanController } from './workbench-plan-controller' import { prepareImageAttachmentUpload } from '../lib/image-attachment-upload' import { isChatAttachmentUploadEnabled } from '../lib/attachment-upload-availability' -import { normalizeWorkspaceRoot } from '../lib/workspace-path' +import { isConversationWorkspacePath, normalizeWorkspaceRoot } from '../lib/workspace-path' import { useKeyboardShortcutSettings } from '../lib/keyboard-shortcut-settings' import { collectComposerChangeSummary } from '../lib/composer-change-summary' import { formatWorkspacePickerError } from '../lib/format-workspace-picker-error' @@ -273,6 +279,15 @@ function isPdfAttachmentFile(file: File): boolean { return file.type === 'application/pdf' || file.name.toLowerCase().endsWith('.pdf') } +function arrayBufferToBase64(buffer: ArrayBuffer): string { + const bytes = new Uint8Array(buffer) + let binary = '' + for (let index = 0; index < bytes.length; index += 0x8000) { + binary += String.fromCharCode(...bytes.subarray(index, index + 0x8000)) + } + return btoa(binary) +} + function stripTransientAttachmentFields(attachments: AttachmentReference[]): AttachmentReference[] { return attachments.map(({ documentText: _documentText, ...attachment }) => attachment) } @@ -406,6 +421,7 @@ export function Workbench(): ReactElement { route, pluginHostRoute, workspaceRoot, + conversationWorkspaceRoot, runtimeConnection, setRoute, openCode, @@ -471,6 +487,7 @@ export function Workbench(): ReactElement { route: s.route, pluginHostRoute: s.pluginHostRoute, workspaceRoot: s.workspaceRoot, + conversationWorkspaceRoot: s.conversationWorkspaceRoot, runtimeConnection: s.runtimeConnection, setRoute: s.setRoute, openCode: s.openCode, @@ -1193,20 +1210,32 @@ export function Workbench(): ReactElement { if (!localFilePath || typeof window.kunGui?.readLocalPdfText !== 'function') { throw new Error(t('composerPdfAttachmentUnavailable')) } + if (!attachmentCapabilities || typeof provider.uploadAttachment !== 'function') { + throw new Error(t('composerAttachmentUnavailable')) + } const result = await window.kunGui.readLocalPdfText({ path: localFilePath }) if (!result.ok) throw new Error(result.message) const documentText = result.text.trim() if (!documentText) throw new Error(t('composerPdfAttachmentNoText')) - uploaded.push({ - id: `doc_${result.mtimeMs}_${index}_${file.name || 'pdf'}`, - kind: 'document', + const attachment = await provider.uploadAttachment({ name: file.name || fileNameFromPath(result.path), mimeType: 'application/pdf', - byteSize: result.size, + dataBase64: arrayBufferToBase64(await file.arrayBuffer()), + documentText, pageCount: result.pageCount, - truncated: result.truncated, + localFilePath, + ...(activeThreadId ? { threadId: activeThreadId } : {}), + ...(workspace ? { workspace } : {}) + }) + uploaded.push({ + id: attachment.id, + kind: 'document', + name: attachment.name, + mimeType: attachment.mimeType, + byteSize: attachment.byteSize, + pageCount: attachment.pageCount, + truncated: attachment.truncated, textPreview: documentText.slice(0, 240), - documentText }) continue } @@ -1277,9 +1306,8 @@ export function Workbench(): ReactElement { const sendWritePrompt = (value: string): void => { const v = value.trim() const attachments = composerAttachments - const imageAttachments = attachments.filter((attachment) => attachment.kind !== 'document') const documentAttachments = attachments.filter((attachment) => attachment.kind === 'document') - const attachmentIds = imageAttachments.map((attachment) => attachment.id) + const attachmentIds = attachments.map((attachment) => attachment.id) const publicAttachments = stripTransientAttachmentFields(attachments) if (!v && attachmentIds.length === 0 && documentAttachments.length === 0) return if (attachmentIds.length > 0 && !attachmentUploadEnabled) { @@ -1506,30 +1534,27 @@ export function Workbench(): ReactElement { } const startNewSddRequirement = async (): Promise => { - const activeCodeWorkspace = activeThreadId - ? normalizeWorkspaceRoot(codeThreads.find((thread) => thread.id === activeThreadId)?.workspace ?? '') - : '' - let targetWorkspace = activeCodeWorkspace || normalizeWorkspaceRoot(workspaceRoot) - if (!targetWorkspace) { - const picked = await chooseWorkspace({ selectThreadAfter: false }) - targetWorkspace = normalizeWorkspaceRoot(picked ?? useChatStore.getState().workspaceRoot) + const suggestedWorkspace = resolveSddRequirementWorkspace(codeThreads, activeThreadId, workspaceRoot) + let targetWorkspace = '' + try { + const picked = await window.kunGui.pickWorkspaceDirectory(suggestedWorkspace || undefined) + if (picked.canceled || !picked.path) return + targetWorkspace = normalizeWorkspaceRoot(picked.path) + } catch (error) { + setError(formatWorkspacePickerError(error)) + return } if (!targetWorkspace) { setError(t('workspaceRequiredToCreateThread')) return } - const restored = await restoreRememberedSddDraft({ - workspaceRoot: targetWorkspace, - readWorkspaceFile: window.kunGui.readWorkspaceFile - }) - if (restored.kind === 'restored') { - await openSddRequirementDraft(restored.draft, restored.content, { - lastSavedContent: restored.lastSavedContent, - saveStatus: restored.saveStatus - }) + if (isConversationWorkspacePath(targetWorkspace, conversationWorkspaceRoot)) { + setError(t('workspaceInsideConversationDir')) return } + if (useSddDraftStore.getState().activeDraft && !await saveActiveSddDraftToDisk()) return + const draftUuid = globalThis.crypto?.randomUUID?.() ?? `draft-${Date.now()}` const draft = createSddDraft({ id: draftUuid, workspaceRoot: targetWorkspace }) const initialContent = [ @@ -1642,9 +1667,8 @@ export function Workbench(): ReactElement { const v = value.trim() const draft = useSddDraftStore.getState().activeDraft const attachments = composerAttachments - const imageAttachments = attachments.filter((attachment) => attachment.kind !== 'document') const documentAttachments = attachments.filter((attachment) => attachment.kind === 'document') - const attachmentIds = imageAttachments.map((attachment) => attachment.id) + const attachmentIds = attachments.map((attachment) => attachment.id) const publicAttachments = stripTransientAttachmentFields(attachments) if ((!v && attachmentIds.length === 0 && documentAttachments.length === 0) || !draft) return if (attachmentIds.length > 0 && !attachmentUploadEnabled) { @@ -2027,9 +2051,8 @@ export function Workbench(): ReactElement { const handleSendAsync = async (): Promise => { const v = input.trim() const attachments = route === 'chat' || route === 'write' ? composerAttachments : [] - const imageAttachments = attachments.filter((attachment) => attachment.kind !== 'document') const documentAttachments = attachments.filter((attachment) => attachment.kind === 'document') - const attachmentIds = imageAttachments.map((attachment) => attachment.id) + const attachmentIds = attachments.map((attachment) => attachment.id) const publicAttachments = stripTransientAttachmentFields(attachments) const fileReferences = route === 'chat' ? composerFileReferences : [] const userFileReferences = composerReferencesToUserFileReferences(fileReferences) @@ -2404,6 +2427,8 @@ export function Workbench(): ReactElement { composerProviderId={resolvedWriteAssistantProviderId} composerPickList={writeAssistantPickList} composerModelGroups={composerModelGroups} + skillCommands={runtimeSkills} + disabledSkillIds={disabledSkillIds} composerReasoningEffort={composerReasoningEffort} setComposerModel={setWriteAssistantModel} setComposerReasoningEffort={setComposerReasoningEffort} diff --git a/src/renderer/src/components/chat/FloatingComposer.tsx b/src/renderer/src/components/chat/FloatingComposer.tsx index 52a566f6d..07a81dcc7 100644 --- a/src/renderer/src/components/chat/FloatingComposer.tsx +++ b/src/renderer/src/components/chat/FloatingComposer.tsx @@ -82,7 +82,7 @@ import { formatCompactNumber, formatCost, formatPercent, - primaryCacheHitRate, + cumulativeCacheHitRate, useThreadUsageState } from '../../hooks/use-thread-usage' import { buildContextCapacity, estimateBlockTokens } from '../../lib/context-capacity' @@ -2485,7 +2485,7 @@ export function FloatingComposer({ · {t('sessionUsageCache', { - cache: formatPercent(primaryCacheHitRate(threadUsage)) + cache: formatPercent(cumulativeCacheHitRate(threadUsage)) })} diff --git a/src/renderer/src/components/chat/GitBranchPicker.tsx b/src/renderer/src/components/chat/GitBranchPicker.tsx index e58a96cd3..c9f0fccb2 100644 --- a/src/renderer/src/components/chat/GitBranchPicker.tsx +++ b/src/renderer/src/components/chat/GitBranchPicker.tsx @@ -2,8 +2,16 @@ import { useCallback, useEffect, useMemo, useRef, useState, type ReactElement } import { createPortal } from 'react-dom' import { AlertCircle, Check, ChevronDown, GitBranch, GitFork, Loader2, Plus, Search } from 'lucide-react' import { useTranslation } from 'react-i18next' +import { + applyGitBranchPrefix, + DEFAULT_GIT_BRANCH_PREFIX, + normalizeGitBranchPrefix, + type AppSettingsV1 +} from '@shared/app-settings' import type { GitBranchesResult, GitBranchRow } from '@shared/git-branches' import { getProvider } from '../../agent/registry' +import { rendererRuntimeClient } from '../../agent/runtime-client' +import { SETTINGS_CHANGED_EVENT } from '../../lib/keyboard-shortcut-settings' import { middleEllipsize } from '../../lib/middle-ellipsize' import { forgetThreadWorktree, @@ -46,6 +54,7 @@ export function GitBranchPicker({ workspaceRoot }: Props): ReactElement | null { const [actingKind, setActingKind] = useState<'switch' | 'worktree' | null>(null) const [error, setError] = useState(null) const [tooltip, setTooltip] = useState(null) + const [branchPrefix, setBranchPrefix] = useState(DEFAULT_GIT_BRANCH_PREFIX) const wrapRef = useRef(null) const inputRef = useRef(null) @@ -77,6 +86,22 @@ export function GitBranchPicker({ workspaceRoot }: Props): ReactElement | null { void load() }, [load]) + useEffect(() => { + let cancelled = false + const apply = (settings: Pick): void => { + if (!cancelled) setBranchPrefix(normalizeGitBranchPrefix(settings.gitBranchPrefix)) + } + void rendererRuntimeClient.getSettings().then(apply).catch(() => undefined) + const onSettingsChanged = (event: Event): void => { + apply((event as CustomEvent).detail) + } + window.addEventListener(SETTINGS_CHANGED_EVENT, onSettingsChanged) + return () => { + cancelled = true + window.removeEventListener(SETTINGS_CHANGED_EVENT, onSettingsChanged) + } + }, []) + useEffect(() => { if (!open) return void load() @@ -106,18 +131,19 @@ export function GitBranchPicker({ workspaceRoot }: Props): ReactElement | null { }, [branches, query]) const trimmedQuery = query.trim() - const exactBranchExists = branches.some((branch) => branch.name === trimmedQuery) - const switchTargetRow = exactBranchExists - ? branches.find((branch) => branch.name === trimmedQuery) ?? null - : null - const canCreate = trimmedQuery.length > 0 && !exactBranchExists + const createBranchName = applyGitBranchPrefix(trimmedQuery, branchPrefix) + const switchTargetRow = branches.find((branch) => branch.name === trimmedQuery) + ?? branches.find((branch) => branch.name === createBranchName) + ?? null + const canCreate = createBranchName.length > 0 && !switchTargetRow const canCreateWorktree = trimmedQuery.length > 0 && (canCreate || Boolean(switchTargetRow)) const currentBranch = result?.ok ? result.currentBranch : null const label = currentBranch || (result?.ok ? t('gitDetached') : t('gitBranchUnavailable')) - const footerBranchLabel = middleEllipsize(trimmedQuery, BRANCH_FOOTER_LABEL_MAX_LENGTH) + const footerBranchLabel = middleEllipsize(createBranchName, BRANCH_FOOTER_LABEL_MAX_LENGTH) const footerCreateLabel = t('gitCreateNamedBranch', { branch: footerBranchLabel }) - const footerCreateTitle = t('gitCreateNamedBranch', { branch: trimmedQuery }) - const footerWorktreeTitle = t('gitNewBranchWorktree', { branch: trimmedQuery }) + const footerCreateTitle = t('gitCreateNamedBranch', { branch: createBranchName }) + const footerWorktreeBranch = canCreate ? createBranchName : switchTargetRow?.name ?? trimmedQuery + const footerWorktreeTitle = t('gitNewBranchWorktree', { branch: footerWorktreeBranch }) const showTooltip = useCallback((text: string, clientX: number, clientY: number): void => { if (!text.trim()) return setTooltip({ text, ...branchTooltipPosition(clientX, clientY) }) @@ -230,7 +256,7 @@ export function GitBranchPicker({ workspaceRoot }: Props): ReactElement | null { } const createAndSwitchBranch = async (): Promise => { - const branch = query.trim() + const branch = createBranchName if (!root || !branch) return setActingBranch(branch) setActingKind('switch') @@ -280,7 +306,7 @@ export function GitBranchPicker({ workspaceRoot }: Props): ReactElement | null { } const createBranchWorktree = async (): Promise => { - const branch = query.trim() + const branch = createBranchName if (!root || !branch) return setActingBranch(branch) setActingKind('worktree') @@ -469,7 +495,7 @@ export function GitBranchPicker({ workspaceRoot }: Props): ReactElement | null { void createAndSwitchBranch() }} > - {actingBranch === trimmedQuery && actingKind === 'switch' ? ( + {actingBranch === createBranchName && actingKind === 'switch' ? ( ) : ( @@ -495,11 +521,11 @@ export function GitBranchPicker({ workspaceRoot }: Props): ReactElement | null { if (canCreate) { void createBranchWorktree() } else { - void checkoutBranchWorktree(trimmedQuery) + void checkoutBranchWorktree(footerWorktreeBranch) } }} > - {actingBranch === trimmedQuery && actingKind === 'worktree' ? ( + {actingBranch === footerWorktreeBranch && actingKind === 'worktree' ? ( ) : ( diff --git a/src/renderer/src/components/chat/MessageTimeline.tsx b/src/renderer/src/components/chat/MessageTimeline.tsx index 791542d81..b0e7f13e0 100644 --- a/src/renderer/src/components/chat/MessageTimeline.tsx +++ b/src/renderer/src/components/chat/MessageTimeline.tsx @@ -70,6 +70,19 @@ export function liveTurnProgressClass(hasActiveGoal: boolean): string { : 'flex w-fit max-w-full items-center gap-2 py-0.5 text-[14px] font-medium text-ds-muted' } +export function activeTimelineTurnKey( + positions: readonly { key: string; top: number }[], + threshold = 96 +): string | null { + if (positions.length === 0) return null + let active = positions[0].key + for (const position of positions) { + if (position.top > threshold) break + active = position.key + } + return active +} + function blockScrollStamp(block: ChatBlock | undefined): string { if (!block) return '' switch (block.kind) { @@ -184,6 +197,7 @@ export function MessageTimeline({ const endRef = useRef(null) const containerRef = useRef(null) const turnRefMap = useRef(new Map()) + const [activeTurnKey, setActiveTurnKey] = useState(null) const turns = useMemo(() => groupTurns(blocks), [blocks]) const latestBlock = blocks[blocks.length - 1] @@ -247,6 +261,39 @@ export function MessageTimeline({ ? Math.max(0, activeThread.forkedFromTurnCount) : undefined + useEffect(() => { + const container = containerRef.current + if (!container || visibleTurnAnchors.length === 0) { + setActiveTurnKey(null) + return + } + let frame: number | null = null + const update = (): void => { + frame = null + if (container.scrollHeight - container.scrollTop - container.clientHeight <= 2) { + setActiveTurnKey(visibleTurnAnchors.at(-1)?.key ?? null) + return + } + const containerTop = container.getBoundingClientRect().top + const positions = visibleTurnAnchors.flatMap((anchor) => { + const node = turnRefMap.current.get(anchor.key) + return node ? [{ key: anchor.key, top: node.getBoundingClientRect().top - containerTop }] : [] + }) + setActiveTurnKey(activeTimelineTurnKey(positions)) + } + const schedule = (): void => { + if (frame === null) frame = window.requestAnimationFrame(update) + } + container.addEventListener('scroll', schedule, { passive: true }) + window.addEventListener('resize', schedule) + schedule() + return () => { + container.removeEventListener('scroll', schedule) + window.removeEventListener('resize', schedule) + if (frame !== null) window.cancelAnimationFrame(frame) + } + }, [visibleTurnAnchors]) + // Tick a clock while a turn is running so the live "Worked for Xs" updates. const [tickNow, setTickNow] = useState(() => Date.now()) useEffect(() => { @@ -259,6 +306,7 @@ export function MessageTimeline({ const jumpToTurn = (key: string): void => { const target = turnRefMap.current.get(key) if (!target) return + setActiveTurnKey(key) target.scrollIntoView({ behavior: 'smooth', block: 'start' }) } @@ -274,9 +322,10 @@ export function MessageTimeline({