diff --git a/CHANGELOG.md b/CHANGELOG.md new file mode 100644 index 0000000000..fe3c4d50b9 --- /dev/null +++ b/CHANGELOG.md @@ -0,0 +1,25 @@ +# Changelog + +## Unreleased + +### Hardening phases 1-5 + +This change set collects the first five maintenance hardening phases from the +Rive deep-read follow-up work. + +| Phase | Area | Summary | +| --- | --- | --- | +| 1 | Runtime permission and usage handling | Made stream watchdog pause/resume accounting robust for concurrent tool calls, added permission timeout handling, integrated Office document abort propagation, and fixed cache/reasoning token usage extraction. | +| 2 | Session JSONL recovery | Recovered sessions with corrupt JSONL rows by parsing message lines independently, surfacing landed corrupt rows as `system_note`, and dropping malformed truncated tail rows. | +| 3 | Bot and OpenGateway abuse controls | Added bot inbound rate and session-binding limits, bounded bot dedupe state, forced bot-bound sessions to `explore`, and capped OpenGateway SSE connections with idle cleanup. | +| 4 | Credential-store secret kind expansion | Extended encrypted credential-store support for bot tokens, bot app secrets, proxy passwords, gateway tokens, and Tavily API keys while preserving legacy API-key/OAuth-token key formats. | +| 5 | Connection credential IPC input hardening | Added shared main-process validation for renderer-controlled connection slugs and API keys before store, credential, or provider side effects. | + +### Verification + +- Runtime package typecheck/build and focused runtime tests. +- Storage package build and focused session-store tests. +- Desktop main build/typecheck and focused bot/OpenGateway, credential-store, + settings/web-search, connection IPC, OAuth, and model-provider regression + suites. +- `git diff --check` before each pushed phase. diff --git a/apps/desktop/src/main/__tests__/bot-incoming-idempotency-contract.test.ts b/apps/desktop/src/main/__tests__/bot-incoming-idempotency-contract.test.ts index 38b1e8a457..e9de21e7b2 100644 --- a/apps/desktop/src/main/__tests__/bot-incoming-idempotency-contract.test.ts +++ b/apps/desktop/src/main/__tests__/bot-incoming-idempotency-contract.test.ts @@ -29,4 +29,103 @@ describe('Bot incoming idempotency contract (PR-BOT-INCOMING-IDEMPOTENCY-0)', () assert.match(main, /const BOT_RECENT_SOURCE_EVENT_LIMIT = 1_000;/, 'dedupe set must stay bounded'); assert.match(main, /while \(botRecentSourceEventKeys\.size > BOT_RECENT_SOURCE_EVENT_LIMIT\)/, 'dedupe set must evict old entries'); }); + + it('expires source-event dedupe entries by TTL before checking duplicates', async () => { + const main = await readRepo('apps/desktop/src/main/main.ts'); + const remember = main.match(/function rememberBotSourceEvent\([^)]*\): boolean \{[\s\S]*?\n\}/); + const prune = main.match(/function pruneExpiredBotSourceEvents\([^)]*\): void \{[\s\S]*?\n\}/); + + assert.ok(remember, 'rememberBotSourceEvent must exist'); + assert.ok(prune, 'pruneExpiredBotSourceEvents must exist'); + assert.match(main, /const BOT_RECENT_SOURCE_EVENT_TTL_MS = 60 \* 60 \* 1_000;/); + assert.ok( + remember![0].indexOf('pruneExpiredBotSourceEvents(now)') < remember![0].indexOf('botRecentSourceEventKeys.has(key)'), + 'expired dedupe entries must be pruned before duplicate lookup', + ); + assert.match( + prune![0], + /if \(now - seenAt <= BOT_RECENT_SOURCE_EVENT_TTL_MS\) break;[\s\S]*botRecentSourceEventKeys\.delete\(key\)/, + 'dedupe TTL cleanup should delete expired oldest entries and keep the hard cap as secondary protection', + ); + }); + + it('rate-limits and session-caps bot turns before create/send side effects', async () => { + const main = await readRepo('apps/desktop/src/main/main.ts'); + const processBlock = main.match(/async function processBotIncomingMessage\([^)]*\): Promise \{[\s\S]*?\n\}\n\nasync function collectBotReply/); + + assert.ok(processBlock, 'processBotIncomingMessage block must exist'); + assert.match(main, /const BOT_CONVERSATION_SESSION_LIMIT = 500;/); + assert.match(main, /const BOT_CONVERSATION_RATE_BURST = 8;/); + assert.match(main, /const BOT_CONVERSATION_RATE_REFILL_MS = 5_000;/); + assert.match(main, /const BOT_CONVERSATION_RATE_BUCKET_TTL_MS = 60 \* 60 \* 1_000;/); + assert.match(main, /const BOT_CONVERSATION_RATE_BUCKET_LIMIT = 1_000;/); + const consumeToken = main.match(/function consumeBotConversationToken\([^]*?\n\}\n\nasync function sendTransientBotNotice/); + const pruneBuckets = main.match(/function pruneExpiredBotConversationRateBuckets\([^)]*\): void \{[\s\S]*?\n\}/); + assert.ok(consumeToken, 'consumeBotConversationToken helper must exist'); + assert.ok(pruneBuckets, 'rate bucket TTL pruning helper must exist'); + assert.match(consumeToken![0], /BOT_CONVERSATION_RATE_BURST/); + assert.match(consumeToken![0], /BOT_CONVERSATION_RATE_REFILL_MS/); + assert.match(consumeToken![0], /pruneExpiredBotConversationRateBuckets\(now\)/); + assert.match(consumeToken![0], /bucket\.tokens -= 1/); + assert.match(consumeToken![0], /while \(botConversationRateBuckets\.size > BOT_CONVERSATION_RATE_BUCKET_LIMIT\)/); + assert.match( + pruneBuckets![0], + /now - bucket\.updatedAt > BOT_CONVERSATION_RATE_BUCKET_TTL_MS[\s\S]*botConversationRateBuckets\.delete\(key\)/, + 'stale rate buckets must be pruned independently of the session cap', + ); + assert.match(main, /async function sendTransientBotNotice[\s\S]*ephemeralTtlMs: ttlMs/); + + const block = processBlock![0]; + const newSessionBranch = block.match(/if \(!sessionId\) \{[\s\S]*?const ready = await getReadyConnection/); + assert.ok(newSessionBranch, 'new bot conversation branch must exist'); + assert.ok( + block.indexOf('consumeBotConversationToken(conversationKey)') < block.indexOf('runtime.createSession'), + 'rate limit must run before creating a bot session', + ); + assert.ok( + block.indexOf('consumeBotConversationToken(conversationKey)') < block.indexOf('runtime.sendMessage'), + 'rate limit must run before runtime.sendMessage', + ); + assert.ok( + block.indexOf('botConversationSessions.size >= BOT_CONVERSATION_SESSION_LIMIT') < block.indexOf('runtime.createSession'), + 'new bot binding cap must run before creating the 501st session', + ); + assert.ok( + newSessionBranch![0].indexOf('botConversationSessions.size >= BOT_CONVERSATION_SESSION_LIMIT') + < newSessionBranch![0].indexOf('consumeBotConversationToken(conversationKey)'), + 'new conversations rejected by a full binding cap must not allocate rate buckets', + ); + assert.match( + block, + /if \(!consumeBotConversationToken\(conversationKey\)\) \{[\s\S]*sendTransientBotNotice[\s\S]*return;/, + 'rate-limited turns must send at most a transient notice and return', + ); + assert.match( + block, + /if \(botConversationSessions\.size >= BOT_CONVERSATION_SESSION_LIMIT\) \{[\s\S]*sendTransientBotNotice[\s\S]*return;/, + 'session-cap rejections must send at most a transient notice and return', + ); + }); + + it('forces existing bot-bound sessions back to explore before send or refuses the turn', async () => { + const main = await readRepo('apps/desktop/src/main/main.ts'); + const processBlock = main.match(/async function processBotIncomingMessage\([^)]*\): Promise \{[\s\S]*?\n\}\n\nasync function collectBotReply/); + const guard = main.match(/async function ensureBotSessionExploreMode\([^)]*\): Promise \{[\s\S]*?\n\}/); + + assert.ok(processBlock, 'processBotIncomingMessage block must exist'); + assert.ok(guard, 'ensureBotSessionExploreMode guard must exist'); + const block = processBlock![0]; + assert.ok( + block.indexOf('ensureBotSessionExploreMode(sessionId, message, SYSTEM_NOTICE_TTL_MS)') < block.indexOf('ensureSessionCanSend(sessionId)'), + 'existing bot sessions must be permission-checked before generic send readiness', + ); + assert.ok( + block.indexOf('ensureBotSessionExploreMode(sessionId, message, SYSTEM_NOTICE_TTL_MS)') < block.indexOf('runtime.sendMessage'), + 'existing bot sessions must be forced/refused before runtime.sendMessage', + ); + assert.match(guard![0], /const header = await store\.readHeader\(sessionId\)/); + assert.match(guard![0], /if \(header\.permissionMode === 'explore'\) return true;/); + assert.match(guard![0], /await runtime\.updateSession\(sessionId, \{ permissionMode: 'explore' \}\);[\s\S]*return true;/); + assert.match(guard![0], /catch \{[\s\S]*sendTransientBotNotice[\s\S]*return false;/); + }); }); diff --git a/apps/desktop/src/main/__tests__/connection-credential-ipc-hardening-contract.test.ts b/apps/desktop/src/main/__tests__/connection-credential-ipc-hardening-contract.test.ts new file mode 100644 index 0000000000..f56100e196 --- /dev/null +++ b/apps/desktop/src/main/__tests__/connection-credential-ipc-hardening-contract.test.ts @@ -0,0 +1,167 @@ +import assert from 'node:assert/strict'; +import { readFileSync } from 'node:fs'; +import { join } from 'node:path'; +import { describe, it } from 'node:test'; + +const mainSource = readFileSync(join(process.cwd(), 'src/main/main.ts'), 'utf8'); + +function handlerBlock(channel: string): string { + const start = mainSource.indexOf(`ipcMain.handle('${channel}'`); + assert.notEqual(start, -1, `${channel} handler must exist`); + const next = mainSource.indexOf('ipcMain.handle(', start + 1); + return mainSource.slice(start, next === -1 ? undefined : next); +} + +describe('connection credential IPC hardening contract', () => { + it('defines shared fail-closed slug and apiKey IPC validators', () => { + assert.match(mainSource, /const IPC_CONNECTION_SLUG_MAX_LENGTH = 64;/); + assert.match(mainSource, /const IPC_CONNECTION_SECRET_MAX_LENGTH = 4096;/); + assert.match(mainSource, /const IPC_CONTROL_CHARACTER_PATTERN = \/\[\\u0000-\\u001F\\u007F\]\//); + assert.match(mainSource, /const IPC_CONNECTION_SLUG_PATTERN = \/\^\[A-Za-z0-9\._-\]\+\$\//); + assert.match( + mainSource, + /function normalizeConnectionSlugForIpc\(value: unknown, label: string\): string \{[\s\S]*typeof value !== 'string'[\s\S]*value\.length === 0[\s\S]*value\.length > IPC_CONNECTION_SLUG_MAX_LENGTH[\s\S]*!IPC_CONNECTION_SLUG_PATTERN\.test\(value\) \|\| IPC_CONTROL_CHARACTER_PATTERN\.test\(value\)[\s\S]*return value;/, + ); + assert.match( + mainSource, + /function normalizeConnectionApiKeyForIpc\(value: unknown, label: string\): string \{[\s\S]*typeof value !== 'string'[\s\S]*value\.length > IPC_CONNECTION_SECRET_MAX_LENGTH[\s\S]*IPC_CONTROL_CHARACTER_PATTERN\.test\(value\)[\s\S]*return value;/, + ); + }); + + it('rejects unsafe slug classes while preserving representative existing valid slugs', () => { + assert.match(mainSource, /const IPC_CONNECTION_SLUG_MAX_LENGTH = 64;/); + assert.match(mainSource, /const IPC_CONNECTION_SLUG_PATTERN = \/\^\[A-Za-z0-9\._-\]\+\$\//); + assert.match(mainSource, /const IPC_CONTROL_CHARACTER_PATTERN = \/\[\\u0000-\\u001F\\u007F\]\//); + assert.match(mainSource, /value\.length === 0/); + assert.match(mainSource, /value\.length > IPC_CONNECTION_SLUG_MAX_LENGTH/); + assert.match(mainSource, /IPC_CONTROL_CHARACTER_PATTERN\.test\(value\)/); + + const helper = mainSource.match(/function normalizeConnectionSlugForIpc\(value: unknown, label: string\): string \{[\s\S]*?\n\}/)?.[0] ?? ''; + assert.ok(helper, 'normalizeConnectionSlugForIpc helper must exist'); + assert.match(helper, /IPC_CONNECTION_SLUG_PATTERN\.test\(value\)/, 'slug validator must reject whitespace, path separators, and colon via an allowlist'); + assert.match( + helper, + /(^|[^\w])\.\.(?!\.)|includes\('\\.\\.'\)|includes\("\.\."\)|traversal|path traversal/i, + 'slug validator must explicitly reject traversal-looking ".." slugs even though dots are otherwise allowed for compatibility', + ); + for (const validSlug of ['claude-subscription', 'codex-subscription', 'zai-coding-plan', 'env-openai']) { + assert.doesNotMatch(validSlug, /[\u0000-\u001F\u007F/:\\]/, `${validSlug} should stay representative-valid`); + assert.ok(validSlug.length <= 64, `${validSlug} should stay under the IPC slug cap`); + } + }); + + it('normalizes create slug and apiKey before store or credential writes', () => { + const helper = mainSource.match(/function normalizeCreateConnectionInput\(input: CreateConnectionInput\): CreateConnectionInput \{[\s\S]*?\n\}/)?.[0] ?? ''; + assert.match(helper, /normalizeConnectionApiKeyForIpc\(input\.apiKey, 'apiKey'\)/); + assert.match(helper, /normalizeConnectionSlugForIpc\(input\.slug, 'connection slug'\)/); + + const handler = handlerBlock('connections:create'); + assert.match(handler, /const normalizedInput = normalizeCreateConnectionInput\(input\);[\s\S]*connectionStore\.create\(normalizedInput\);/); + assert.match(handler, /credentialStore\.setSecret\(connection\.slug, 'api_key', normalizedInput\.apiKey\)/); + }); + + it('caps and validates create apiKey before persistence without echoing the secret value', () => { + const apiKeyHelper = mainSource.match(/function normalizeConnectionApiKeyForIpc\(value: unknown, label: string\): string \{[\s\S]*?\n\}/)?.[0] ?? ''; + const createHelper = mainSource.match(/function normalizeCreateConnectionInput\(input: CreateConnectionInput\): CreateConnectionInput \{[\s\S]*?\n\}/)?.[0] ?? ''; + const handler = handlerBlock('connections:create'); + + assert.match(apiKeyHelper, /value\.length > IPC_CONNECTION_SECRET_MAX_LENGTH/); + assert.match(apiKeyHelper, /IPC_CONTROL_CHARACTER_PATTERN\.test\(value\)/); + assert.doesNotMatch(apiKeyHelper, /String\(value\)|\$\{value\}|value\}/, 'apiKey validation errors must not echo cleartext secret values'); + assert.ok( + createHelper.indexOf('normalizeConnectionApiKeyForIpc(input.apiKey, \'apiKey\')') + < createHelper.indexOf('normalizeConnectionSlugForIpc(input.slug, \'connection slug\')'), + 'create must validate apiKey before constructing normalized input for persistence', + ); + assert.ok( + handler.indexOf('normalizeCreateConnectionInput(input)') + < handler.indexOf('connectionStore.create(normalizedInput)'), + 'create must normalize and cap apiKey before connection persistence', + ); + assert.ok( + handler.indexOf('normalizeCreateConnectionInput(input)') + < handler.indexOf('credentialStore.setSecret'), + 'create must normalize and cap apiKey before credential persistence', + ); + }); + + it('normalizes update slug and apiKey before side-effecting update or credential writes', () => { + const helper = mainSource.match(/async function normalizeUpdateConnectionInput\([\s\S]*?\n\}/)?.[0] ?? ''; + assert.match(helper, /const normalizedPatch = normalizeConnectionPatchSecretsForIpc\(patch\);[\s\S]*const existing = await connectionStore\.get\(slug\);/); + + const handler = handlerBlock('connections:update'); + assert.match(handler, /slug = normalizeConnectionSlugForIpc\(slug, 'connection slug'\);/); + assert.match(handler, /const normalizedPatch = await normalizeUpdateConnectionInput\(slug, patch\);/); + assert.match(handler, /connectionStore\.update\(slug, normalizedPatch\)/); + assert.match(handler, /credentialStore\.(?:setSecret|deleteSecret)\(slug, 'api_key'/); + }); + + it('validates renderer-controlled slug handlers before store, credential, or provider work', () => { + for (const channel of [ + 'connections:setDefault', + 'connections:delete', + 'connections:test', + 'connections:fetchModels', + ]) { + const handler = handlerBlock(channel); + const normalizeAt = handler.indexOf('normalizeConnectionSlugForIpc'); + assert.notEqual(normalizeAt, -1, `${channel} must normalize slug`); + for (const sideEffect of [ + 'connectionStore.', + 'credentialStore.', + 'resolveConnectionSecret(', + 'testConnection(', + 'fetchProviderModels(', + ]) { + const sideEffectAt = handler.indexOf(sideEffect); + if (sideEffectAt !== -1) { + assert.ok(normalizeAt < sideEffectAt, `${channel} must normalize before ${sideEffect}`); + } + } + } + + const hasSecret = handlerBlock('connections:hasSecret'); + assert.match( + hasSecret, + /slug = normalizeConnectionSlugForIpc\(slug, 'connection slug'\);[\s\S]*return Boolean\(await resolveConnectionSecret\(slug\)\);/, + ); + }); + + it('preserves update apiKey clearing semantics while rejecting invalid provided strings', () => { + const helper = mainSource.match(/function normalizeConnectionPatchSecretsForIpc\(patch: UpdateConnectionInput\): UpdateConnectionInput \{[\s\S]*?\n\}/)?.[0] ?? ''; + assert.match(helper, /if \(!Object\.prototype\.hasOwnProperty\.call\(patch, 'apiKey'\)\) return patch;/); + assert.match(helper, /if \(patch\.apiKey === undefined\) return patch;/); + assert.match(helper, /apiKey: normalizeConnectionApiKeyForIpc\(patch\.apiKey, 'apiKey'\)/); + + const handler = handlerBlock('connections:update'); + assert.match(handler, /if \(normalizedPatch\.apiKey !== undefined\) \{[\s\S]*if \(normalizedPatch\.apiKey\) await credentialStore\.setSecret\(slug, 'api_key', normalizedPatch\.apiKey\);[\s\S]*else await credentialStore\.deleteSecret\(slug, 'api_key'\);/); + }); + + it('keeps OAuth baseUrl normalization and provider-aware secret resolution wired', () => { + assert.match( + mainSource, + /function normalizeCreateConnectionInput\(input: CreateConnectionInput\): CreateConnectionInput \{[\s\S]*defaults\.authKind === 'oauth_token'[\s\S]*baseUrl: defaults\.baseUrl/, + 'create must continue forcing canonical OAuth provider baseUrl', + ); + assert.match( + mainSource, + /async function normalizeUpdateConnectionInput\([\s\S]*PROVIDER_DEFAULTS\[providerType\]\.authKind === 'oauth_token'[\s\S]*baseUrl: PROVIDER_DEFAULTS\[providerType\]\.baseUrl/, + 'update must continue forcing canonical OAuth provider baseUrl', + ); + assert.match(mainSource, /connections:test[\s\S]*const apiKey = await resolveConnectionSecret\(slug\)/); + assert.match(mainSource, /connections:fetchModels[\s\S]*const apiKey = await resolveConnectionSecret\(slug\)/); + assert.match(mainSource, /connections:hasSecret[\s\S]*return Boolean\(await resolveConnectionSecret\(slug\)\)/); + }); + + it('does not echo cleartext API keys in thrown errors or IPC return values', () => { + const validatorRegion = mainSource.match(/function normalizeConnectionApiKeyForIpc[\s\S]*?function normalizeCreateConnectionInput/)?.[0] ?? ''; + const createHandler = handlerBlock('connections:create'); + const updateHandler = handlerBlock('connections:update'); + + for (const source of [validatorRegion, createHandler, updateHandler]) { + assert.doesNotMatch(source, /errorMessage:\s*[^,\n]*apiKey/); + assert.doesNotMatch(source, /throw new Error\([^)]*apiKey[^)]*\$\{[^}]*value/); + assert.doesNotMatch(source, /return \{[\s\S]*apiKey[\s\S]*\}/); + } + }); +}); diff --git a/apps/desktop/src/main/__tests__/credential-store-contract.test.ts b/apps/desktop/src/main/__tests__/credential-store-contract.test.ts new file mode 100644 index 0000000000..5dc4bfcf4a --- /dev/null +++ b/apps/desktop/src/main/__tests__/credential-store-contract.test.ts @@ -0,0 +1,58 @@ +import assert from 'node:assert/strict'; +import { readFileSync } from 'node:fs'; +import { join } from 'node:path'; +import { describe, it } from 'node:test'; + +const source = readFileSync(join(process.cwd(), 'src/main/credential-store.ts'), 'utf8'); + +describe('credential store secret-kind expansion contract', () => { + it('keeps legacy connection secret key names backward compatible', () => { + assert.match(source, /case 'api_key':\s*return 'apiKey';/); + assert.match(source, /case 'oauth_token':\s*return 'oauthToken';/); + }); + + it('declares the Phase 1 settings secret kinds', () => { + for (const kind of [ + 'bot_token', + 'app_secret', + 'proxy_password', + 'gateway_token', + 'tavily_api_key', + ]) { + assert.match(source, new RegExp(`'${kind}'`), `${kind} must be a public CredentialKind`); + } + }); + + it('exposes typed helpers for future settings migration without changing settings consumers', () => { + for (const method of [ + 'getBotToken', + 'setBotToken', + 'deleteBotToken', + 'getBotAppSecret', + 'setBotAppSecret', + 'deleteBotAppSecret', + 'getProxyPassword', + 'setProxyPassword', + 'deleteProxyPassword', + 'getGatewayToken', + 'setGatewayToken', + 'deleteGatewayToken', + 'getTavilyApiKey', + 'setTavilyApiKey', + 'deleteTavilyApiKey', + ]) { + assert.match(source, new RegExp(`${method}\\(`), `${method} must be present`); + } + }); + + it('uses deterministic non-secret slugs for provider-scoped and global settings secrets', () => { + assert.match(source, /const BOT_SECRET_SLUG_PREFIX = 'settings:bot';/); + assert.match(source, /function botSecretSlug\(provider: BotProvider\): string \{\s*return `\$\{BOT_SECRET_SLUG_PREFIX\}:\$\{provider\}`;\s*\}/); + assert.match(source, /const GLOBAL_PROXY_SECRET_SLUG = 'settings:network-proxy';/); + assert.match(source, /const GLOBAL_GATEWAY_SECRET_SLUG = 'settings:open-gateway';/); + assert.match(source, /const GLOBAL_TAVILY_SECRET_SLUG = 'settings:web-search:tavily';/); + assert.doesNotMatch(source, /token\}:/, 'raw bot tokens must not be interpolated into key names'); + assert.doesNotMatch(source, /password\}:/, 'raw proxy passwords must not be interpolated into key names'); + assert.doesNotMatch(source, /secret\}:/, 'raw app secrets must not be interpolated into key names'); + }); +}); diff --git a/apps/desktop/src/main/__tests__/credential-store-secret-kinds-contract.test.ts b/apps/desktop/src/main/__tests__/credential-store-secret-kinds-contract.test.ts new file mode 100644 index 0000000000..5baf452349 --- /dev/null +++ b/apps/desktop/src/main/__tests__/credential-store-secret-kinds-contract.test.ts @@ -0,0 +1,143 @@ +import assert from 'node:assert/strict'; +import { readFile } from 'node:fs/promises'; +import { join, resolve } from 'node:path'; +import { describe, it } from 'node:test'; +import type { BotProvider } from '@maka/core'; +import type { CredentialKind, CredentialStore } from '../credential-store.js'; + +const repoRoot = process.cwd().endsWith(join('apps', 'desktop')) + ? resolve(process.cwd(), '..', '..') + : process.cwd(); + +async function readRepo(relativePath: string): Promise { + return readFile(join(repoRoot, relativePath), 'utf8'); +} + +const credentialKinds: CredentialKind[] = [ + 'api_key', + 'oauth_token', + 'bot_token', + 'app_secret', + 'proxy_password', + 'gateway_token', + 'tavily_api_key', +]; + +type Phase4CredentialStoreMethods = Pick< + CredentialStore, + | 'getBotToken' + | 'setBotToken' + | 'deleteBotToken' + | 'getBotAppSecret' + | 'setBotAppSecret' + | 'deleteBotAppSecret' + | 'getProxyPassword' + | 'setProxyPassword' + | 'deleteProxyPassword' + | 'getGatewayToken' + | 'setGatewayToken' + | 'deleteGatewayToken' + | 'getTavilyApiKey' + | 'setTavilyApiKey' + | 'deleteTavilyApiKey' +>; + +type BotScopedSetter = (provider: BotProvider, secret: string) => Promise; +const botScopedSetter: BotScopedSetter = null as unknown as Phase4CredentialStoreMethods['setBotToken']; + +void credentialKinds; +void botScopedSetter; + +describe('credential-store secret kind expansion contract', () => { + it('exposes the Phase 4 credential kinds without changing legacy stored kind names', async () => { + const source = await readRepo('apps/desktop/src/main/credential-store.ts'); + + for (const kind of credentialKinds) { + assert.match(source, new RegExp(`'${kind}'`), `CredentialKind must include ${kind}`); + } + + assert.match(source, /case 'api_key':\s*return 'apiKey';/, 'api_key must keep the legacy stored key suffix'); + assert.match(source, /case 'oauth_token':\s*return 'oauthToken';/, 'oauth_token must keep the legacy stored key suffix'); + assert.match(source, /getApiKey\(slug: string\)[\s\S]*?return this\.get\(slug, 'apiKey'\);/); + assert.match(source, /setApiKey\(slug: string, apiKey: string\)[\s\S]*?return this\.set\(slug, 'apiKey', apiKey\);/); + assert.match(source, /getOAuthToken\(slug: string\)[\s\S]*?return this\.get\(slug, 'oauthToken'\);/); + assert.match(source, /setOAuthToken\(slug: string, token: string\)[\s\S]*?return this\.set\(slug, 'oauthToken', token\);/); + assert.match(source, /private key\(slug: string, kind: StoredCredentialKind\): string \{\s*return `\$\{slug\}:\$\{kind\}`;\s*\}/); + }); + + it('scopes bot token and app-secret helpers by provider and kind', async () => { + const source = await readRepo('apps/desktop/src/main/credential-store.ts'); + + assert.match(source, /const BOT_SECRET_SLUG_PREFIX = 'settings:bot';/); + assert.match(source, /function botSecretSlug\(provider: BotProvider\): string \{\s*return `\$\{BOT_SECRET_SLUG_PREFIX\}:\$\{provider\}`;\s*\}/); + assert.match(source, /getBotToken\(provider: BotProvider\)[\s\S]*return this\.get\(botSecretSlug\(provider\), 'botToken'\);/); + assert.match(source, /setBotToken\(provider: BotProvider, token: string\)[\s\S]*return this\.set\(botSecretSlug\(provider\), 'botToken', token\);/); + assert.match(source, /deleteBotToken\(provider: BotProvider\)[\s\S]*return this\.deleteSecret\(botSecretSlug\(provider\), 'bot_token'\);/); + assert.match(source, /getBotAppSecret\(provider: BotProvider\)[\s\S]*return this\.get\(botSecretSlug\(provider\), 'botAppSecret'\);/); + assert.match(source, /setBotAppSecret\(provider: BotProvider, secret: string\)[\s\S]*return this\.set\(botSecretSlug\(provider\), 'botAppSecret', secret\);/); + assert.match(source, /deleteBotAppSecret\(provider: BotProvider\)[\s\S]*return this\.deleteSecret\(botSecretSlug\(provider\), 'app_secret'\);/); + assert.doesNotMatch(source, /botSecretSlug\([^)]*(token|secret|value|apiKey|password|key)[^)]*\)/); + }); + + it('uses isolated singleton slugs for global secret helpers', async () => { + const source = await readRepo('apps/desktop/src/main/credential-store.ts'); + + assert.match(source, /const GLOBAL_PROXY_SECRET_SLUG = 'settings:network-proxy';/); + assert.match(source, /const GLOBAL_GATEWAY_SECRET_SLUG = 'settings:open-gateway';/); + assert.match(source, /const GLOBAL_TAVILY_SECRET_SLUG = 'settings:web-search:tavily';/); + assert.match(source, /getProxyPassword\(\)[\s\S]*return this\.get\(GLOBAL_PROXY_SECRET_SLUG, 'proxyPassword'\);/); + assert.match(source, /setProxyPassword\(password: string\)[\s\S]*return this\.set\(GLOBAL_PROXY_SECRET_SLUG, 'proxyPassword', password\);/); + assert.match(source, /deleteProxyPassword\(\)[\s\S]*return this\.deleteSecret\(GLOBAL_PROXY_SECRET_SLUG, 'proxy_password'\);/); + assert.match(source, /getGatewayToken\(\)[\s\S]*return this\.get\(GLOBAL_GATEWAY_SECRET_SLUG, 'gatewayToken'\);/); + assert.match(source, /setGatewayToken\(token: string\)[\s\S]*return this\.set\(GLOBAL_GATEWAY_SECRET_SLUG, 'gatewayToken', token\);/); + assert.match(source, /deleteGatewayToken\(\)[\s\S]*return this\.deleteSecret\(GLOBAL_GATEWAY_SECRET_SLUG, 'gateway_token'\);/); + assert.match(source, /getTavilyApiKey\(\)[\s\S]*return this\.get\(GLOBAL_TAVILY_SECRET_SLUG, 'tavilyApiKey'\);/); + assert.match(source, /setTavilyApiKey\(key: string\)[\s\S]*return this\.set\(GLOBAL_TAVILY_SECRET_SLUG, 'tavilyApiKey', key\);/); + assert.match(source, /deleteTavilyApiKey\(\)[\s\S]*return this\.deleteSecret\(GLOBAL_TAVILY_SECRET_SLUG, 'tavily_api_key'\);/); + }); + + it('keeps credential writes encrypted, fail-closed, and in the existing file shape', async () => { + const source = await readRepo('apps/desktop/src/main/credential-store.ts'); + const setBlock = source.match(/private async set\([^)]*\): Promise \{[\s\S]*?\n \}/); + const getBlock = source.match(/private async get\([^)]*\): Promise \{[\s\S]*?\n \}/); + const readBlock = source.match(/private async readUnlocked\(\): Promise \{[\s\S]*?\n \}/); + + assert.ok(setBlock, 'set helper must exist'); + assert.ok(getBlock, 'get helper must exist'); + assert.ok(readBlock, 'readUnlocked helper must exist'); + assert.match(setBlock![0], /if \(!safeStorage\.isEncryptionAvailable\(\)\) \{[\s\S]*throw new Error/); + assert.match(setBlock![0], /file\.values\[this\.key\(slug, kind\)\] = safeStorage\.encryptString\(value\)\.toString\('base64'\);/); + assert.match(getBlock![0], /if \(!encrypted\) return null;/); + assert.match(getBlock![0], /safeStorage\.decryptString\(Buffer\.from\(encrypted, 'base64'\)\)/); + assert.match(readBlock![0], /ENOENT'[\s\S]*return \{ values: \{\} \};/); + assert.match(source, /interface CredentialFile \{\s*values: Record;\s*\}/); + assert.match(source, /JSON\.stringify\(file, null, 2\) \+ '\\n'/); + assert.doesNotMatch(source, /writeFile\([^)]*value/); + }); + + it('keeps delete helpers idempotent and targeted', async () => { + const source = await readRepo('apps/desktop/src/main/credential-store.ts'); + const deleteSecretBlock = source.match(/async deleteSecret\([^)]*\): Promise \{[\s\S]*?\n \}/); + const deleteBlock = source.match(/async delete\(slug: string\): Promise \{[\s\S]*?\n \}/); + + assert.ok(deleteSecretBlock, 'deleteSecret helper must exist'); + assert.ok(deleteBlock, 'delete helper must exist'); + assert.match(deleteSecretBlock![0], /delete file\.values\[this\.key\(slug, toStoredKind\(kind\)\)\];/); + assert.doesNotMatch(deleteSecretBlock![0], /if \([^)]*values\[this\.key/, 'targeted delete should remain idempotent when missing'); + assert.match(deleteBlock![0], /for \(const kind of STORED_CREDENTIAL_KINDS\)/); + assert.match(source, /'apiKey'[\s\S]*'oauthToken'[\s\S]*'botToken'[\s\S]*'botAppSecret'[\s\S]*'proxyPassword'[\s\S]*'gatewayToken'[\s\S]*'tavilyApiKey'/); + }); + + it('does not start settings migration or change renderer-facing settings masking in this phase', async () => { + const main = await readRepo('apps/desktop/src/main/main.ts'); + const helpers = await readRepo('apps/desktop/src/main/settings-ipc-helpers.ts'); + + assert.match(main, /ipcMain\.handle\('settings:get', async \(\) => maskAppSettings\(await settingsStore\.get\(\)\)\);/); + assert.doesNotMatch(main, /credentialStore\.(getBotToken|getBotAppSecret|getProxyPassword|getGatewayToken|getTavilyApiKey)/); + assert.match(helpers, /password: shouldReveal\(revealPatch\.network\?\.proxy\?\.password\)/); + assert.match(helpers, /token: shouldReveal\(revealPatch\.botChat\?\.channels\?\.\[provider as BotProvider\]\?\.token\)/); + assert.match(helpers, /appSecret: shouldReveal\(revealPatch\.botChat\?\.channels\?\.\[provider as BotProvider\]\?\.appSecret\)/); + assert.match(helpers, /token: shouldReveal\(revealPatch\.openGateway\?\.token\)/); + assert.match(helpers, /apiKey: maskSensitive\(settings\.webSearch\.providers\.tavily\.apiKey\) \?\? ''/); + }); +}); diff --git a/apps/desktop/src/main/__tests__/office-document-tool.test.ts b/apps/desktop/src/main/__tests__/office-document-tool.test.ts index af3466951a..6f1024928b 100644 --- a/apps/desktop/src/main/__tests__/office-document-tool.test.ts +++ b/apps/desktop/src/main/__tests__/office-document-tool.test.ts @@ -268,6 +268,40 @@ describe('OfficeDocument read-only tool', () => { }); assert.equal(timeout.ok, false); assert.equal(timeout.ok ? null : timeout.reason, 'officecli_timeout'); + + const preAbortedController = new AbortController(); + preAbortedController.abort(); + const preAborted = await runOfficeDocumentOperation({ + cwd: workspaceRoot, + path: 'sheet.xlsx', + operation: 'validate', + abortSignal: preAbortedController.signal, + runner: fakeRunner(() => { + throw new Error('pre-aborted operation should not run officecli'); + }), + }); + assert.equal(preAborted.ok, false); + assert.equal(preAborted.ok ? null : preAborted.reason, 'officecli_aborted'); + + const inflightController = new AbortController(); + const inflight = await runOfficeDocumentOperation({ + cwd: workspaceRoot, + path: 'sheet.xlsx', + operation: 'validate', + abortSignal: inflightController.signal, + runner: fakeRunner((_cmd, _args, options, callback) => { + assert.equal(options.signal, inflightController.signal); + (options.signal as AbortSignal).addEventListener('abort', () => { + const error = new Error('aborted') as NodeJS.ErrnoException; + error.name = 'AbortError'; + error.code = 'ABORT_ERR'; + callback(error, '', ''); + }, { once: true }); + inflightController.abort(); + }), + }); + assert.equal(inflight.ok, false); + assert.equal(inflight.ok ? null : inflight.reason, 'officecli_aborted'); }); }); diff --git a/apps/desktop/src/main/__tests__/open-gateway-sse-abuse-contract.test.ts b/apps/desktop/src/main/__tests__/open-gateway-sse-abuse-contract.test.ts new file mode 100644 index 0000000000..4bfe93fbd7 --- /dev/null +++ b/apps/desktop/src/main/__tests__/open-gateway-sse-abuse-contract.test.ts @@ -0,0 +1,50 @@ +import { strict as assert } from 'node:assert'; +import { readFile } from 'node:fs/promises'; +import { resolve } from 'node:path'; +import { describe, it } from 'node:test'; + +const REPO_ROOT = resolve(import.meta.dirname, '../../../../..'); + +async function readRepo(path: string): Promise { + return readFile(resolve(REPO_ROOT, path), 'utf8'); +} + +describe('OpenGateway SSE abuse hardening contract', () => { + it('keeps stream limits closed and rejects before SSE headers are written', async () => { + const source = await readRepo('apps/desktop/src/main/open-gateway.ts'); + const openStream = source.match(/private openSessionEventStream\([^]*?\n private removeEventClient/); + + assert.ok(openStream, 'openSessionEventStream block must exist'); + assert.match(source, /const OPEN_GATEWAY_EVENT_STREAM_TOTAL_LIMIT = 10;/); + assert.match(source, /const OPEN_GATEWAY_EVENT_STREAM_PER_SESSION_LIMIT = 3;/); + assert.ok( + openStream![0].indexOf('too_many_event_streams') < openStream![0].indexOf("res.setHeader('Content-Type', 'text/event-stream; charset=utf-8')"), + 'limit rejection must happen before event-stream headers are committed', + ); + assert.match( + openStream![0], + /writeJson\(res, 429, \{ ok: false, error: 'too_many_event_streams' \}\);[\s\S]*return;/, + 'excess streams must use the stable closed 429 error shape', + ); + }); + + it('bases idle timeout on real SSE events, not heartbeat comments, and clears timers on removal', async () => { + const source = await readRepo('apps/desktop/src/main/open-gateway.ts'); + const openStream = source.match(/private openSessionEventStream\([^]*?\n private removeEventClient/); + const removeClient = source.match(/private removeEventClient\([^]*?\n private closeEventClients/); + + assert.ok(openStream, 'openSessionEventStream block must exist'); + assert.ok(removeClient, 'removeEventClient block must exist'); + assert.match(source, /const OPEN_GATEWAY_EVENT_IDLE_TIMEOUT_MS = 5 \* 60 \* 1_000;/); + assert.match(openStream![0], /heartbeat: setInterval\(\(\) => \{[\s\S]*res\.write\(`: heartbeat \$\{this\.now\(\)\}\\n\\n`\);[\s\S]*\}, OPEN_GATEWAY_EVENT_HEARTBEAT_MS\)/); + assert.doesNotMatch( + openStream![0].match(/heartbeat: setInterval\(\(\) => \{[\s\S]*?\}, OPEN_GATEWAY_EVENT_HEARTBEAT_MS\)/)?.[0] ?? '', + /resetIdleTimer/, + 'heartbeat comments must not reset the idle timer', + ); + assert.match(openStream![0], /write\(chunk\) \{[\s\S]*res\.write\(chunk\);[\s\S]*resetIdleTimer\(\);[\s\S]*\}/); + assert.match(removeClient![0], /if \(client\.closed\) return;[\s\S]*client\.closed = true;/); + assert.match(removeClient![0], /clearInterval\(client\.heartbeat\);/); + assert.match(removeClient![0], /clearTimeout\(client\.idleTimeout\);/); + }); +}); diff --git a/apps/desktop/src/main/__tests__/open-gateway.test.ts b/apps/desktop/src/main/__tests__/open-gateway.test.ts index 52b2f71cd8..470f230958 100644 --- a/apps/desktop/src/main/__tests__/open-gateway.test.ts +++ b/apps/desktop/src/main/__tests__/open-gateway.test.ts @@ -562,6 +562,64 @@ describe('OpenGatewayService', () => { assert.ok(statusChanges.includes(0), 'closing an SSE stream should publish activeEventStreams=0'); }); + test('rejects excess SSE streams before establishing event-stream headers', async () => { + const service = makeService(); + activeServices.push(service); + const status = await service.sync(createGatewaySettings({ enabled: true, port: 0, token: 'dev-token' }).openGateway); + assert.ok(status.baseUrl); + + const controllers: AbortController[] = []; + try { + for (let index = 0; index < 3; index += 1) { + const opened = await openEventStream(status.baseUrl, 'same-session'); + controllers.push(opened.controller); + assert.equal(opened.response.status, 200); + } + assert.equal(service.getStatus().activeEventStreams, 3); + + const perSessionRejected = await fetchJson(`${status.baseUrl}/v1/sessions/same-session/events`, 'dev-token'); + assert.equal(perSessionRejected.status, 429); + assert.equal(perSessionRejected.body.error, 'too_many_event_streams'); + assert.doesNotMatch(perSessionRejected.headers.get('content-type') ?? '', /^text\/event-stream/); + assert.equal(service.getStatus().activeEventStreams, 3); + + for (let index = 0; index < 7; index += 1) { + const opened = await openEventStream(status.baseUrl, `other-${index}`); + controllers.push(opened.controller); + assert.equal(opened.response.status, 200); + } + assert.equal(service.getStatus().activeEventStreams, 10); + + const globalRejected = await fetchJson(`${status.baseUrl}/v1/sessions/global-overflow/events`, 'dev-token'); + assert.equal(globalRejected.status, 429); + assert.equal(globalRejected.body.error, 'too_many_event_streams'); + assert.doesNotMatch(globalRejected.headers.get('content-type') ?? '', /^text\/event-stream/); + assert.equal(service.getStatus().activeEventStreams, 10); + } finally { + for (const controller of controllers) controller.abort(); + await waitFor(() => service.getStatus().activeEventStreams === 0); + } + }); + + test('stop closes active SSE clients and clears stream counts', async () => { + const service = makeService(); + activeServices.push(service); + const status = await service.sync(createGatewaySettings({ enabled: true, port: 0, token: 'dev-token' }).openGateway); + assert.ok(status.baseUrl); + + const first = await openEventStream(status.baseUrl, 's1'); + const second = await openEventStream(status.baseUrl, 's2'); + assert.equal(service.getStatus().activeEventStreams, 2); + + await service.stop(); + await Promise.all([ + readUntilClosed(first.response.body!.getReader()), + readUntilClosed(second.response.body!.getReader()), + ]); + + assert.equal(service.getStatus().activeEventStreams, 0); + }); + test('replays recent SSE events after Last-Event-ID cursor', async () => { const service = makeService(); activeServices.push(service); @@ -951,6 +1009,18 @@ async function fetchJson( }; } +async function openEventStream( + baseUrl: string, + sessionId: string, +): Promise<{ controller: AbortController; response: Response }> { + const controller = new AbortController(); + const response = await fetch(`${baseUrl}/v1/sessions/${sessionId}/events`, { + headers: { Authorization: 'Bearer dev-token' }, + signal: controller.signal, + }); + return { controller, response }; +} + function session(overrides: Partial & { id: string }): SessionSummary { return { name: overrides.id, diff --git a/apps/desktop/src/main/credential-store.ts b/apps/desktop/src/main/credential-store.ts index 71772ae729..1627d79150 100644 --- a/apps/desktop/src/main/credential-store.ts +++ b/apps/desktop/src/main/credential-store.ts @@ -1,9 +1,24 @@ import { safeStorage } from 'electron'; import { mkdir, readFile, rename, writeFile } from 'node:fs/promises'; import { dirname, join } from 'node:path'; - -type StoredCredentialKind = 'apiKey' | 'oauthToken'; -export type CredentialKind = 'api_key' | 'oauth_token'; +import type { BotProvider } from '@maka/core'; + +type StoredCredentialKind = + | 'apiKey' + | 'oauthToken' + | 'botToken' + | 'botAppSecret' + | 'proxyPassword' + | 'gatewayToken' + | 'tavilyApiKey'; +export type CredentialKind = + | 'api_key' + | 'oauth_token' + | 'bot_token' + | 'app_secret' + | 'proxy_password' + | 'gateway_token' + | 'tavily_api_key'; interface CredentialFile { values: Record; @@ -13,6 +28,21 @@ export interface CredentialStore { getSecret(slug: string, kind: CredentialKind): Promise; setSecret(slug: string, kind: CredentialKind, value: string): Promise; deleteSecret(slug: string, kind?: CredentialKind): Promise; + getBotToken(provider: BotProvider): Promise; + setBotToken(provider: BotProvider, token: string): Promise; + deleteBotToken(provider: BotProvider): Promise; + getBotAppSecret(provider: BotProvider): Promise; + setBotAppSecret(provider: BotProvider, secret: string): Promise; + deleteBotAppSecret(provider: BotProvider): Promise; + getProxyPassword(): Promise; + setProxyPassword(password: string): Promise; + deleteProxyPassword(): Promise; + getGatewayToken(): Promise; + setGatewayToken(token: string): Promise; + deleteGatewayToken(): Promise; + getTavilyApiKey(): Promise; + setTavilyApiKey(key: string): Promise; + deleteTavilyApiKey(): Promise; getApiKey(slug: string): Promise; getOAuthToken(slug: string): Promise; setApiKey(slug: string, apiKey: string): Promise; @@ -49,6 +79,66 @@ class SafeStorageCredentialStore implements CredentialStore { }); } + getBotToken(provider: BotProvider): Promise { + return this.get(botSecretSlug(provider), 'botToken'); + } + + setBotToken(provider: BotProvider, token: string): Promise { + return this.set(botSecretSlug(provider), 'botToken', token); + } + + deleteBotToken(provider: BotProvider): Promise { + return this.deleteSecret(botSecretSlug(provider), 'bot_token'); + } + + getBotAppSecret(provider: BotProvider): Promise { + return this.get(botSecretSlug(provider), 'botAppSecret'); + } + + setBotAppSecret(provider: BotProvider, secret: string): Promise { + return this.set(botSecretSlug(provider), 'botAppSecret', secret); + } + + deleteBotAppSecret(provider: BotProvider): Promise { + return this.deleteSecret(botSecretSlug(provider), 'app_secret'); + } + + getProxyPassword(): Promise { + return this.get(GLOBAL_PROXY_SECRET_SLUG, 'proxyPassword'); + } + + setProxyPassword(password: string): Promise { + return this.set(GLOBAL_PROXY_SECRET_SLUG, 'proxyPassword', password); + } + + deleteProxyPassword(): Promise { + return this.deleteSecret(GLOBAL_PROXY_SECRET_SLUG, 'proxy_password'); + } + + getGatewayToken(): Promise { + return this.get(GLOBAL_GATEWAY_SECRET_SLUG, 'gatewayToken'); + } + + setGatewayToken(token: string): Promise { + return this.set(GLOBAL_GATEWAY_SECRET_SLUG, 'gatewayToken', token); + } + + deleteGatewayToken(): Promise { + return this.deleteSecret(GLOBAL_GATEWAY_SECRET_SLUG, 'gateway_token'); + } + + getTavilyApiKey(): Promise { + return this.get(GLOBAL_TAVILY_SECRET_SLUG, 'tavilyApiKey'); + } + + setTavilyApiKey(key: string): Promise { + return this.set(GLOBAL_TAVILY_SECRET_SLUG, 'tavilyApiKey', key); + } + + deleteTavilyApiKey(): Promise { + return this.deleteSecret(GLOBAL_TAVILY_SECRET_SLUG, 'tavily_api_key'); + } + getApiKey(slug: string): Promise { return this.get(slug, 'apiKey'); } @@ -68,8 +158,9 @@ class SafeStorageCredentialStore implements CredentialStore { async delete(slug: string): Promise { await this.withQueue(async () => { const file = await this.readUnlocked(); - delete file.values[this.key(slug, 'apiKey')]; - delete file.values[this.key(slug, 'oauthToken')]; + for (const kind of STORED_CREDENTIAL_KINDS) { + delete file.values[this.key(slug, kind)]; + } await this.write(file); }); } @@ -118,6 +209,40 @@ class SafeStorageCredentialStore implements CredentialStore { } } +const STORED_CREDENTIAL_KINDS = [ + 'apiKey', + 'oauthToken', + 'botToken', + 'botAppSecret', + 'proxyPassword', + 'gatewayToken', + 'tavilyApiKey', +] as const satisfies readonly StoredCredentialKind[]; + +const BOT_SECRET_SLUG_PREFIX = 'settings:bot'; +const GLOBAL_PROXY_SECRET_SLUG = 'settings:network-proxy'; +const GLOBAL_GATEWAY_SECRET_SLUG = 'settings:open-gateway'; +const GLOBAL_TAVILY_SECRET_SLUG = 'settings:web-search:tavily'; + +function botSecretSlug(provider: BotProvider): string { + return `${BOT_SECRET_SLUG_PREFIX}:${provider}`; +} + function toStoredKind(kind: CredentialKind): StoredCredentialKind { - return kind === 'api_key' ? 'apiKey' : 'oauthToken'; + switch (kind) { + case 'api_key': + return 'apiKey'; + case 'oauth_token': + return 'oauthToken'; + case 'bot_token': + return 'botToken'; + case 'app_secret': + return 'botAppSecret'; + case 'proxy_password': + return 'proxyPassword'; + case 'gateway_token': + return 'gatewayToken'; + case 'tavily_api_key': + return 'tavilyApiKey'; + } } diff --git a/apps/desktop/src/main/main.ts b/apps/desktop/src/main/main.ts index 48db550635..ab94e5a355 100644 --- a/apps/desktop/src/main/main.ts +++ b/apps/desktop/src/main/main.ts @@ -408,34 +408,90 @@ async function resolveConnectionSecret(slug: string): Promise { return credentialStore.getSecret(slug, 'api_key'); } +const IPC_CONNECTION_SLUG_MAX_LENGTH = 64; +const IPC_CONNECTION_SECRET_MAX_LENGTH = 4096; +const IPC_CONTROL_CHARACTER_PATTERN = /[\u0000-\u001F\u007F]/; +const IPC_CONNECTION_SLUG_PATTERN = /^[A-Za-z0-9._-]+$/; + +function hasTraversalLookingSlugSegment(value: string): boolean { + return value.split('.').some((segment) => segment.length === 0); +} + +function normalizeConnectionSlugForIpc(value: unknown, label: string): string { + if (typeof value !== 'string') { + throw new Error(`${label} must be a string`); + } + if (value.length === 0) { + throw new Error(`${label} is required`); + } + if (value.length > IPC_CONNECTION_SLUG_MAX_LENGTH) { + throw new Error(`${label} must be ${IPC_CONNECTION_SLUG_MAX_LENGTH} characters or fewer`); + } + if (!IPC_CONNECTION_SLUG_PATTERN.test(value) || IPC_CONTROL_CHARACTER_PATTERN.test(value)) { + throw new Error(`${label} contains invalid characters`); + } + if (hasTraversalLookingSlugSegment(value)) { + throw new Error(`${label} contains invalid path traversal segments`); + } + return value; +} + +function normalizeConnectionApiKeyForIpc(value: unknown, label: string): string { + if (typeof value !== 'string') { + throw new Error(`${label} must be a string`); + } + if (value.length > IPC_CONNECTION_SECRET_MAX_LENGTH) { + throw new Error(`${label} must be ${IPC_CONNECTION_SECRET_MAX_LENGTH} characters or fewer`); + } + if (IPC_CONTROL_CHARACTER_PATTERN.test(value)) { + throw new Error(`${label} contains invalid characters`); + } + return value; +} + function normalizeCreateConnectionInput(input: CreateConnectionInput): CreateConnectionInput { - const defaults = PROVIDER_DEFAULTS[input.providerType]; + const apiKey = input.apiKey === undefined + ? undefined + : normalizeConnectionApiKeyForIpc(input.apiKey, 'apiKey'); + const slug = normalizeConnectionSlugForIpc(input.slug, 'connection slug'); + const normalizedInput = { ...input, slug, ...(apiKey !== undefined ? { apiKey } : {}) }; + const defaults = PROVIDER_DEFAULTS[normalizedInput.providerType]; if (defaults.authKind === 'oauth_token') { - return { ...input, baseUrl: defaults.baseUrl }; + return { ...normalizedInput, baseUrl: defaults.baseUrl }; } - if (input.baseUrl === undefined) return input; - const result = normalizeConnectionBaseUrl(input.baseUrl); + if (normalizedInput.baseUrl === undefined) return normalizedInput; + const result = normalizeConnectionBaseUrl(normalizedInput.baseUrl); if (!result.ok) { throw new Error(result.error); } - return { ...input, baseUrl: result.value }; + return { ...normalizedInput, baseUrl: result.value }; +} + +function normalizeConnectionPatchSecretsForIpc(patch: UpdateConnectionInput): UpdateConnectionInput { + if (!Object.prototype.hasOwnProperty.call(patch, 'apiKey')) return patch; + if (patch.apiKey === undefined) return patch; + return { + ...patch, + apiKey: normalizeConnectionApiKeyForIpc(patch.apiKey, 'apiKey'), + }; } async function normalizeUpdateConnectionInput( slug: string, patch: UpdateConnectionInput, ): Promise { + const normalizedPatch = normalizeConnectionPatchSecretsForIpc(patch); const existing = await connectionStore.get(slug); const providerType = existing?.providerType; if (providerType && PROVIDER_DEFAULTS[providerType].authKind === 'oauth_token') { - return { ...patch, baseUrl: PROVIDER_DEFAULTS[providerType].baseUrl }; + return { ...normalizedPatch, baseUrl: PROVIDER_DEFAULTS[providerType].baseUrl }; } - if (patch.baseUrl === undefined) return patch; - const result = normalizeConnectionBaseUrl(patch.baseUrl); + if (normalizedPatch.baseUrl === undefined) return normalizedPatch; + const result = normalizeConnectionBaseUrl(normalizedPatch.baseUrl); if (!result.ok) { throw new Error(result.error); } - return { ...patch, baseUrl: result.value }; + return { ...normalizedPatch, baseUrl: result.value }; } const planReminderStore = createPlanReminderStore(workspaceRoot); @@ -894,7 +950,19 @@ const runtime = new SessionManager({ const botConversationSessions = new Map(); const botConversationQueues = new Map>(); const botRecentSourceEventKeys = new Map(); +const botConversationRateBuckets = new Map(); const BOT_RECENT_SOURCE_EVENT_LIMIT = 1_000; +const BOT_RECENT_SOURCE_EVENT_TTL_MS = 60 * 60 * 1_000; +const BOT_CONVERSATION_SESSION_LIMIT = 500; +const BOT_CONVERSATION_RATE_BURST = 8; +const BOT_CONVERSATION_RATE_REFILL_MS = 5_000; +const BOT_CONVERSATION_RATE_BUCKET_TTL_MS = 60 * 60 * 1_000; +const BOT_CONVERSATION_RATE_BUCKET_LIMIT = 1_000; + +interface BotConversationRateBucket { + tokens: number; + updatedAt: number; +} // PR110b: onboarding service composes existing stores + runtime to // derive `OnboardingState` and manage `OnboardingMilestone[]`. @@ -2314,10 +2382,11 @@ function registerIpc(): void { }); ipcMain.handle('connections:getDefault', () => connectionStore.getDefault()); ipcMain.handle('connections:setDefault', async (_event, slug: string | null) => { - if (slug && !(await connectionStore.get(slug))) { - throw new Error(`No such connection: ${slug}`); + const normalizedSlug = slug === null ? null : normalizeConnectionSlugForIpc(slug, 'connection slug'); + if (normalizedSlug && !(await connectionStore.get(normalizedSlug))) { + throw new Error(`No such connection: ${normalizedSlug}`); } - await connectionStore.setDefault(slug); + await connectionStore.setDefault(normalizedSlug); emitConnectionListChanged(); }); ipcMain.handle('connections:create', async (_event, input: CreateConnectionInput) => { @@ -2362,6 +2431,7 @@ function registerIpc(): void { // Same OAuth-boundary rule as create: if the current/new provider // uses an OAuth token, force the canonical provider endpoint and // ignore renderer-provided baseUrl text entirely. + slug = normalizeConnectionSlugForIpc(slug, 'connection slug'); const normalizedPatch = await normalizeUpdateConnectionInput(slug, patch); const connection = await connectionStore.update(slug, normalizedPatch); if (normalizedPatch.apiKey !== undefined) { @@ -2372,11 +2442,13 @@ function registerIpc(): void { return connection; }); ipcMain.handle('connections:delete', async (_event, slug: string) => { + slug = normalizeConnectionSlugForIpc(slug, 'connection slug'); await connectionStore.delete(slug); await credentialStore.deleteSecret(slug); emitConnectionListChanged(); }); ipcMain.handle('connections:test', async (_event, slug: string, opts?: { model?: string }) => { + slug = normalizeConnectionSlugForIpc(slug, 'connection slug'); const connection = await connectionStore.get(slug); if (!connection) return { ok: false, errorMessage: `找不到模型连接:${slug}` }; const apiKey = await resolveConnectionSecret(slug); @@ -2395,6 +2467,7 @@ function registerIpc(): void { return result; }); ipcMain.handle('connections:fetchModels', async (_event, slug: string) => { + slug = normalizeConnectionSlugForIpc(slug, 'connection slug'); const connection = await connectionStore.get(slug); if (!connection) throw new Error(`找不到模型连接:${slug}`); const apiKey = await resolveConnectionSecret(slug); @@ -2421,9 +2494,10 @@ function registerIpc(): void { throw new Error(generalizedErrorMessageChinese(error, '拉取模型列表失败')); } }); - ipcMain.handle('connections:hasSecret', async (_event, slug: string) => - Boolean(await resolveConnectionSecret(slug)), - ); + ipcMain.handle('connections:hasSecret', async (_event, slug: string) => { + slug = normalizeConnectionSlugForIpc(slug, 'connection slug'); + return Boolean(await resolveConnectionSecret(slug)); + }); // PR110b: Onboarding snapshot + milestone IPCs. Renderer polls via // these on app load and whenever `sessions:changed` / @@ -2808,8 +2882,10 @@ async function handleBotIncomingMessage(message: BotIncomingMessage): Promise BOT_RECENT_SOURCE_EVENT_LIMIT) { const oldest = botRecentSourceEventKeys.keys().next().value; if (!oldest) break; @@ -2818,6 +2894,59 @@ function rememberBotSourceEvent(message: BotIncomingMessage): boolean { return false; } +function pruneExpiredBotSourceEvents(now: number): void { + for (const [key, seenAt] of botRecentSourceEventKeys) { + if (now - seenAt <= BOT_RECENT_SOURCE_EVENT_TTL_MS) break; + botRecentSourceEventKeys.delete(key); + } +} + +function consumeBotConversationToken(conversationKey: string, now = Date.now()): boolean { + pruneExpiredBotConversationRateBuckets(now); + const bucket = botConversationRateBuckets.get(conversationKey) ?? { + tokens: BOT_CONVERSATION_RATE_BURST, + updatedAt: now, + }; + const elapsed = Math.max(0, now - bucket.updatedAt); + const refilled = Math.floor(elapsed / BOT_CONVERSATION_RATE_REFILL_MS); + if (refilled > 0) { + bucket.tokens = Math.min(BOT_CONVERSATION_RATE_BURST, bucket.tokens + refilled); + bucket.updatedAt += refilled * BOT_CONVERSATION_RATE_REFILL_MS; + } + if (bucket.tokens <= 0) { + botConversationRateBuckets.set(conversationKey, bucket); + return false; + } + bucket.tokens -= 1; + botConversationRateBuckets.set(conversationKey, bucket); + while (botConversationRateBuckets.size > BOT_CONVERSATION_RATE_BUCKET_LIMIT) { + const oldest = botConversationRateBuckets.keys().next().value; + if (!oldest) break; + botConversationRateBuckets.delete(oldest); + } + return true; +} + +function pruneExpiredBotConversationRateBuckets(now: number): void { + for (const [key, bucket] of botConversationRateBuckets) { + if (now - bucket.updatedAt > BOT_CONVERSATION_RATE_BUCKET_TTL_MS) { + botConversationRateBuckets.delete(key); + } + } +} + +async function sendTransientBotNotice(message: BotIncomingMessage, text: string, ttlMs: number): Promise { + await botRegistry.sendMessage( + message.platform, + message.chatId, + text, + { + ...(message.sourceMessageId ? { replyToMessageId: message.sourceMessageId } : {}), + ephemeralTtlMs: ttlMs, + }, + ).catch(() => null); +} + async function processBotIncomingMessage( conversationKey: string, message: BotIncomingMessage, @@ -2852,6 +2981,7 @@ async function processBotIncomingMessage( // member would otherwise be able to wipe everyone else's context. if (isPlaintextResetCommand({ text, isGroup: message.isGroup })) { const had = botConversationSessions.delete(conversationKey); + botConversationRateBuckets.delete(conversationKey); const replyOptions = { ...(message.sourceMessageId ? { replyToMessageId: message.sourceMessageId } : {}), ephemeralTtlMs: SYSTEM_NOTICE_TTL_MS, @@ -2865,6 +2995,22 @@ async function processBotIncomingMessage( let sessionId = botConversationSessions.get(conversationKey); try { if (!sessionId) { + if (botConversationSessions.size >= BOT_CONVERSATION_SESSION_LIMIT) { + await sendTransientBotNotice( + message, + 'Maka 当前机器人会话数量已达上限,请重置或清理旧会话后再试。', + SYSTEM_NOTICE_TTL_MS, + ); + return; + } + if (!consumeBotConversationToken(conversationKey)) { + await sendTransientBotNotice( + message, + 'Maka 收到的机器人消息过于频繁,请稍后再试。', + SYSTEM_NOTICE_TTL_MS, + ); + return; + } const ready = await getReadyConnection(await connectionStore.getDefault(), undefined); const summary = await runtime.createSession({ cwd: process.cwd(), @@ -2881,7 +3027,17 @@ async function processBotIncomingMessage( botConversationSessions.set(conversationKey, sessionId); emitSessionsChanged('created', sessionId); } else { + const permissionModeOk = await ensureBotSessionExploreMode(sessionId, message, SYSTEM_NOTICE_TTL_MS); + if (!permissionModeOk) return; await ensureSessionCanSend(sessionId); + if (!consumeBotConversationToken(conversationKey)) { + await sendTransientBotNotice( + message, + 'Maka 收到的机器人消息过于频繁,请稍后再试。', + SYSTEM_NOTICE_TTL_MS, + ); + return; + } } const turnId = randomUUID(); @@ -2960,6 +3116,27 @@ async function processBotIncomingMessage( } } +async function ensureBotSessionExploreMode( + sessionId: string, + message: BotIncomingMessage, + noticeTtlMs: number, +): Promise { + const header = await store.readHeader(sessionId); + if (header.permissionMode === 'explore') return true; + try { + await runtime.updateSession(sessionId, { permissionMode: 'explore' }); + emitSessionsChanged('updated', sessionId); + return true; + } catch { + await sendTransientBotNotice( + message, + 'Maka 已拒绝这条机器人消息:绑定会话当前不是只读探索模式,请先在桌面端切回 explore 后再试。', + noticeTtlMs, + ); + return false; + } +} + async function collectBotReply( sessionId: string, iterator: AsyncIterable, diff --git a/apps/desktop/src/main/office-document-tool.ts b/apps/desktop/src/main/office-document-tool.ts index 8905d79eff..23a819ee06 100644 --- a/apps/desktop/src/main/office-document-tool.ts +++ b/apps/desktop/src/main/office-document-tool.ts @@ -56,6 +56,7 @@ export type OfficeDocumentResult = | 'invalid_props' | 'file_exists' | 'officecli_missing' + | 'officecli_aborted' | 'officecli_timeout' | 'officecli_failed'; message: string; @@ -98,7 +99,7 @@ export function buildOfficeDocumentTool(): MakaTool< .describe('Optional depth for get; capped at 6.'), }), permissionRequired: false, - impl: async ({ path, operation, topic, viewMode, selector, query, depth }, { cwd }) => runOfficeDocumentOperation({ + impl: async ({ path, operation, topic, viewMode, selector, query, depth }, { cwd, abortSignal }) => runOfficeDocumentOperation({ cwd, path, operation, @@ -107,6 +108,7 @@ export function buildOfficeDocumentTool(): MakaTool< selector, query, depth, + abortSignal, }), }; } @@ -145,7 +147,7 @@ export function buildOfficeDocumentEditTool(): MakaTool< }), permissionRequired: true, categoryHint: 'file_write', - impl: async ({ path, operation, target, elementType, props, index }, { cwd }) => runOfficeDocumentEditOperation({ + impl: async ({ path, operation, target, elementType, props, index }, { cwd, abortSignal }) => runOfficeDocumentEditOperation({ cwd, path, operation, @@ -153,6 +155,7 @@ export function buildOfficeDocumentEditTool(): MakaTool< elementType, props, index, + abortSignal, }), }; } @@ -168,6 +171,7 @@ export async function runOfficeDocumentOperation(input: { depth?: unknown; runner?: OfficeCliRunner; timeoutMs?: number; + abortSignal?: AbortSignal; }): Promise { const operation = normalizeOperation(input.operation); if (!operation) { @@ -188,6 +192,7 @@ export async function runOfficeDocumentOperation(input: { args: buildOfficeHelpArgs(input.topic), runner: input.runner, timeoutMs: input.timeoutMs, + abortSignal: input.abortSignal, }); } @@ -231,6 +236,7 @@ export async function runOfficeDocumentOperation(input: { args: argsResult.args, runner, timeoutMs, + abortSignal: input.abortSignal, }); } @@ -242,12 +248,13 @@ async function runOfficeCliOperation(input: { args: string[]; runner?: OfficeCliRunner; timeoutMs?: number; + abortSignal?: AbortSignal; }): Promise { const workspaceRoot = await realpath(input.cwd); const runner = input.runner ?? execFile; const timeoutMs = input.timeoutMs ?? OFFICE_DOCUMENT_TIMEOUT_MS; try { - const output = await runOfficeCli(runner, input.args, timeoutMs); + const output = await runOfficeCli(runner, input.args, timeoutMs, input.abortSignal); const stdout = sanitizeOfficeCliOutput(output.stdout, workspaceRoot); const stderr = sanitizeOfficeCliOutput(output.stderr, workspaceRoot); const cappedStdout = capOutput(stdout); @@ -276,6 +283,17 @@ async function runOfficeCliOperation(input: { message: '本机未检测到 officecli。请先安装 officecli,并确认 `officecli --version` 可运行后重试。', }; } + if (code === 'ABORT_ERR' || (error as Error).name === 'AbortError') { + return { + kind: 'office_document', + ok: false, + operation: input.operation, + ...(input.relPath ? { path: input.relPath } : {}), + args: input.absPath && input.relPath ? displayArgs(input.args, input.absPath, input.relPath) : input.args, + reason: 'officecli_aborted', + message: 'officecli 操作已取消。', + }; + } if (code === 'ETIMEDOUT' || killed) { return { kind: 'office_document', @@ -309,6 +327,7 @@ export async function runOfficeDocumentEditOperation(input: { index?: unknown; runner?: OfficeCliRunner; timeoutMs?: number; + abortSignal?: AbortSignal; }): Promise { const operation = normalizeEditOperation(input.operation); if (!operation) { @@ -360,6 +379,7 @@ export async function runOfficeDocumentEditOperation(input: { args: argsResult.args, runner: input.runner, timeoutMs: input.timeoutMs, + abortSignal: input.abortSignal, }); } @@ -584,8 +604,17 @@ function normalizeBoundedText(value: unknown): string | null { return text; } -function runOfficeCli(runner: OfficeCliRunner, args: string[], timeoutMs: number): Promise<{ stdout: string; stderr: string }> { +function runOfficeCli( + runner: OfficeCliRunner, + args: string[], + timeoutMs: number, + abortSignal?: AbortSignal, +): Promise<{ stdout: string; stderr: string }> { return new Promise((resolvePromise, reject) => { + if (abortSignal?.aborted) { + reject(abortError()); + return; + } const child = runner( 'officecli', args, @@ -593,6 +622,7 @@ function runOfficeCli(runner: OfficeCliRunner, args: string[], timeoutMs: number timeout: timeoutMs, maxBuffer: OFFICE_DOCUMENT_MAX_BUFFER, env: buildOfficeCliEnv(), + ...(abortSignal ? { signal: abortSignal } : {}), }, (error, stdout, stderr) => { if (error) { @@ -607,6 +637,13 @@ function runOfficeCli(runner: OfficeCliRunner, args: string[], timeoutMs: number }); } +function abortError(): Error { + const error = new Error('officecli aborted') as Error & { code?: string }; + error.name = 'AbortError'; + error.code = 'ABORT_ERR'; + return error; +} + function sanitizeOfficeCliOutput(text: string, workspaceRoot: string): string { return redactSecrets(text.replaceAll(workspaceRoot, '')).trim(); } diff --git a/apps/desktop/src/main/open-gateway.ts b/apps/desktop/src/main/open-gateway.ts index 20caac3ddc..379eb7119b 100644 --- a/apps/desktop/src/main/open-gateway.ts +++ b/apps/desktop/src/main/open-gateway.ts @@ -511,6 +511,14 @@ export class OpenGatewayService { req: IncomingMessage, res: ServerResponse, ): void { + if ( + this.countEventClients() >= OPEN_GATEWAY_EVENT_STREAM_TOTAL_LIMIT || + (this.eventClients.get(sessionId)?.size ?? 0) >= OPEN_GATEWAY_EVENT_STREAM_PER_SESSION_LIMIT + ) { + writeJson(res, 429, { ok: false, error: 'too_many_event_streams' }); + return; + } + res.statusCode = 200; res.setHeader('Content-Type', 'text/event-stream; charset=utf-8'); res.setHeader('Cache-Control', 'no-cache, no-transform'); @@ -519,14 +527,26 @@ export class OpenGatewayService { res.write('retry: 1000\n'); res.write(`: session ${sessionId} connected\n\n`); - const client: GatewayEventClient = { + let client: GatewayEventClient; + const resetIdleTimer = () => { + clearTimeout(client.idleTimeout); + client.idleTimeout = setTimeout(() => { + this.removeEventClient(sessionId, client); + }, OPEN_GATEWAY_EVENT_IDLE_TIMEOUT_MS); + }; + client = { response: res, heartbeat: setInterval(() => { res.write(`: heartbeat ${this.now()}\n\n`); }, OPEN_GATEWAY_EVENT_HEARTBEAT_MS), write(chunk) { res.write(chunk); + resetIdleTimer(); }, + idleTimeout: setTimeout(() => { + this.removeEventClient(sessionId, client); + }, OPEN_GATEWAY_EVENT_IDLE_TIMEOUT_MS), + closed: false, }; const clients = this.eventClients.get(sessionId) ?? new Set(); clients.add(client); @@ -538,7 +558,10 @@ export class OpenGatewayService { } private removeEventClient(sessionId: string, client: GatewayEventClient): void { + if (client.closed) return; + client.closed = true; clearInterval(client.heartbeat); + clearTimeout(client.idleTimeout); const clients = this.eventClients.get(sessionId); if (clients) { clients.delete(client); @@ -664,6 +687,9 @@ const OPEN_GATEWAY_MAX_BODY_BYTES = 16 * 1024; const OPEN_GATEWAY_MESSAGE_PAGE_DEFAULT_LIMIT = 100; const OPEN_GATEWAY_MESSAGE_PAGE_MAX_LIMIT = 200; const OPEN_GATEWAY_EVENT_HEARTBEAT_MS = 15_000; +const OPEN_GATEWAY_EVENT_IDLE_TIMEOUT_MS = 5 * 60 * 1_000; +const OPEN_GATEWAY_EVENT_STREAM_TOTAL_LIMIT = 10; +const OPEN_GATEWAY_EVENT_STREAM_PER_SESSION_LIMIT = 3; const OPEN_GATEWAY_EVENT_REPLAY_LIMIT = 100; const OPEN_GATEWAY_EVENT_RECENT_LIMIT = 50; const OPEN_GATEWAY_REPLAY_CURSOR_LIMIT = 256; @@ -700,6 +726,8 @@ function capGatewayPath(value: string): string { interface GatewayEventClient { response: ServerResponse; heartbeat: ReturnType; + idleTimeout: ReturnType; + closed: boolean; write(chunk: string): void; } diff --git a/packages/runtime/src/__tests__/ai-sdk-backend.test.ts b/packages/runtime/src/__tests__/ai-sdk-backend.test.ts index b6317c0d46..609aee74ed 100644 --- a/packages/runtime/src/__tests__/ai-sdk-backend.test.ts +++ b/packages/runtime/src/__tests__/ai-sdk-backend.test.ts @@ -1,14 +1,18 @@ import assert from 'node:assert/strict'; import { describe, test } from 'node:test'; +import { MockLanguageModelV3, simulateReadableStream } from 'ai/test'; +import type { LanguageModelV3StreamPart } from '@ai-sdk/provider'; import type { LlmConnection, SessionHeader } from '@maka/core'; import type { SessionEvent } from '@maka/core/events'; import type { ToolResultMessage } from '@maka/core/session'; +import type { LlmCallRecord } from '@maka/core/usage-stats/types'; import { AiSdkBackend, INVALID_TOOL_NAME, MAX_ACTIVE_SUBAGENT_TOOLS_PER_TURN, TOOL_ERROR_RESULT_MAX_CHARS, formatSyntheticToolErrorText, + normalizeAiSdkUsage, repairMakaToolCall, type MakaTool, } from '../ai-sdk-backend.js'; @@ -214,7 +218,200 @@ describe('AiSdkBackend stop', () => { }); }); +describe('AiSdkBackend usage telemetry', () => { + test('normalizes standard LanguageModelUsage detail token fields', () => { + const usage = normalizeAiSdkUsage({ + inputTokens: 100, + outputTokens: 20, + inputTokenDetails: { + cacheReadTokens: 30, + cacheWriteTokens: 10, + }, + outputTokenDetails: { + reasoningTokens: 5, + }, + }); + + assert.deepEqual(usage, { + inputTokens: 100, + outputTokens: 20, + cachedInputTokens: 30, + cacheWriteInputTokens: 10, + reasoningTokens: 5, + totalTokens: 120, + }); + }); + + test('normalizes cache and reasoning tokens to messages, events, and telemetry', async () => { + const messages: unknown[] = []; + const events: SessionEvent[] = []; + const llmRecords: LlmCallRecord[] = []; + const chunks: LanguageModelV3StreamPart[] = [ + { type: 'stream-start', warnings: [] }, + { type: 'text-start', id: 'text-1' }, + { type: 'text-delta', id: 'text-1', delta: 'hello' }, + { type: 'text-end', id: 'text-1' }, + { + type: 'finish', + finishReason: { unified: 'stop', raw: 'stop' }, + usage: { + inputTokens: { + total: 10, + noCache: 5, + cacheRead: 3, + cacheWrite: 2, + }, + outputTokens: { + total: 7, + text: 5, + reasoning: 2, + }, + }, + }, + ]; + const model = new MockLanguageModelV3({ + doStream: { + stream: simulateReadableStream({ + chunks, + initialDelayInMs: null, + chunkDelayInMs: null, + }), + }, + }); + const backend = new AiSdkBackend({ + sessionId: 'session-1', + header: header(), + appendMessage: async (message) => { + messages.push(message); + }, + connection: connection(), + apiKey: 'sk-test', + modelId: 'mock-model-id', + permissionEngine: new PermissionEngine({ newId: () => 'permission-id', now: () => 1 }), + modelFactory: () => model, + tools: [], + newId: idGenerator(), + now: monotonicClock(), + recordLlmCall: (record) => { + llmRecords.push(record); + }, + }); + + for await (const event of backend.send({ turnId: 'turn-1', text: 'hi', context: [] })) { + events.push(event); + } + + const usageMessage = messages.find((message) => + (message as { type?: string }).type === 'token_usage' + ) as { input?: number; output?: number; cacheRead?: number; cacheCreation?: number } | undefined; + const usageEvent = events.find((event) => event.type === 'token_usage') as + | Extract + | undefined; + + assert.equal((usageMessage as { type?: string } | undefined)?.type, 'token_usage'); + assert.equal((usageMessage as { turnId?: string } | undefined)?.turnId, 'turn-1'); + assert.equal(usageMessage?.input, 10); + assert.equal(usageMessage?.output, 7); + assert.equal(usageMessage?.cacheRead, 3); + assert.equal(usageMessage?.cacheCreation, 2); + assert.equal(usageEvent?.input, 10); + assert.equal(usageEvent?.output, 7); + assert.equal(usageEvent?.cacheRead, 3); + assert.equal(usageEvent?.cacheCreation, 2); + assert.equal(llmRecords[0]?.inputTokens, 10); + assert.equal(llmRecords[0]?.outputTokens, 7); + assert.equal(llmRecords[0]?.cachedInputTokens, 3); + assert.equal(llmRecords[0]?.cacheWriteInputTokens, 2); + assert.equal(llmRecords[0]?.reasoningTokens, 2); + assert.equal(llmRecords[0]?.totalTokens, 17); + }); +}); + describe('AiSdkBackend tool permission category hints', () => { + test('permission prompt timeout expires one request, resumes watchdog, and writes an error result', async () => { + const messages: unknown[] = []; + const events: SessionEvent[] = []; + const permissionEngine = new PermissionEngine({ newId: idGenerator(), now: () => 1 }); + let implCalled = false; + let pauseCount = 0; + let resumeCount = 0; + const backend = new AiSdkBackend({ + sessionId: 'session-1', + header: header('ask'), + appendMessage: async (message) => { + messages.push(message); + }, + connection: connection(), + apiKey: 'sk-test', + modelId: 'claude-sonnet-4-5-20250929', + permissionEngine, + modelFactory: () => ({}), + tools: [], + newId: idGenerator(), + now: () => 1, + permissionTimeoutMs: 1, + }); + const tool: MakaTool = { + name: 'Write', + description: 'write file', + parameters: {}, + permissionRequired: true, + impl: async () => { + implCalled = true; + return { ok: true }; + }, + }; + (backend as unknown as { + currentWatchdog: { pause(): void; resume(): void }; + }).currentWatchdog = { + pause: () => { + pauseCount += 1; + }, + resume: () => { + resumeCount += 1; + }, + }; + + const execute = (backend as unknown as { + wrapToolExecute( + tool: MakaTool, + turnId: string, + queue: { push(event: SessionEvent): void }, + ): (args: unknown, ctx: { toolCallId: string; abortSignal: AbortSignal }) => Promise; + }).wrapToolExecute(tool, 'turn-1', { push: (event) => events.push(event) }); + + const result = await execute( + { path: 'notes.md', content: 'hello' }, + { toolCallId: 'tool-1', abortSignal: new AbortController().signal }, + ); + const permissionRequest = events.find((event) => event.type === 'permission_request') as + | Extract + | undefined; + const toolResult = events.find((event) => event.type === 'tool_result') as + | Extract + | undefined; + + assert.equal(implCalled, false); + assert.equal(pauseCount, 1); + assert.equal(resumeCount, 1); + assert.equal(permissionEngine.pendingCount('turn-1'), 0); + assert.equal(permissionEngine.recordResponse('turn-1', { + requestId: permissionRequest?.requestId ?? 'missing', + decision: 'allow', + }), null); + assert.match((result as { error?: string }).error ?? '', /Permission flow aborted/); + assert.match((result as { error?: string }).error ?? '', /timed out/); + assert.equal(toolResult?.isError, true); + assert.equal( + messages.some((message) => + (message as { type?: string; toolUseId?: string; isError?: boolean }).type === 'tool_result' && + (message as { toolUseId?: string }).toolUseId === 'tool-1' && + (message as { isError?: boolean }).isError === true, + ), + true, + ); + }); + test('passes categoryHint through PermissionEngine before tool execution', async () => { const messages: unknown[] = []; const events: SessionEvent[] = []; @@ -410,6 +607,62 @@ describe('AiSdkBackend tool permission category hints', () => { { status: 'aborted', toolCallId: 'tool-aborted' }, ]); }); + + test('maps aborted OfficeDocument results to aborted tool telemetry', async () => { + const events: SessionEvent[] = []; + const telemetry: Array<{ status: string; toolCallId?: string }> = []; + const backend = new AiSdkBackend({ + sessionId: 'session-1', + header: header('ask'), + appendMessage: async () => {}, + connection: connection(), + apiKey: 'sk-test', + modelId: 'claude-sonnet-4-5-20250929', + permissionEngine: new PermissionEngine({ newId: () => 'permission-id', now: () => 1 }), + modelFactory: () => ({}), + tools: [], + newId: idGenerator(), + now: () => 1, + recordToolInvocation: (record) => { + telemetry.push({ status: record.status, toolCallId: record.toolCallId }); + }, + }); + const tool: MakaTool = { + name: 'OfficeDocument', + description: 'read office', + parameters: {}, + permissionRequired: false, + impl: async () => ({ + kind: 'office_document', + ok: false, + operation: 'view', + path: 'slides.pptx', + args: ['view', 'slides.pptx', 'outline'], + reason: 'officecli_aborted', + message: 'officecli 操作已取消。', + }), + }; + const execute = (backend as unknown as { + wrapToolExecute( + tool: MakaTool, + turnId: string, + queue: { push(event: SessionEvent): void }, + ): (args: unknown, ctx: { toolCallId: string; abortSignal: AbortSignal }) => Promise; + }).wrapToolExecute(tool, 'turn-1', { push: (event) => events.push(event) }); + + await execute({ path: 'slides.pptx', operation: 'view' }, { + toolCallId: 'tool-office-aborted', + abortSignal: new AbortController().signal, + }); + + assert.equal( + (events.find((event) => event.type === 'tool_result') as { isError?: boolean } | undefined)?.isError, + true, + ); + assert.deepEqual(telemetry, [ + { status: 'aborted', toolCallId: 'tool-office-aborted' }, + ]); + }); }); describe('AiSdkBackend tool-call repair', () => { @@ -500,3 +753,8 @@ function idGenerator(): () => string { let index = 0; return () => `id-${++index}`; } + +function monotonicClock(): () => number { + let value = 1_000; + return () => ++value; +} diff --git a/packages/runtime/src/__tests__/permission-engine.test.ts b/packages/runtime/src/__tests__/permission-engine.test.ts index f05099dacb..f7c5d6f609 100644 --- a/packages/runtime/src/__tests__/permission-engine.test.ts +++ b/packages/runtime/src/__tests__/permission-engine.test.ts @@ -228,6 +228,31 @@ describe('PermissionEngine — turn lifecycle', () => { expect(engine.pendingCount('t1')).toBe(0); }); + test('expireRequest rejects one parked request and ignores late responses', async () => { + const { engine } = makeEngine(); + engine.beginTurn('t1'); + const r = engine.evaluate({ + sessionId: 's1', + turnId: 't1', + toolUseId: 'tu1', + toolName: 'Write', + args: {}, + mode: 'ask', + }); + if (r.kind !== 'prompt') throw new Error('expected prompt'); + const parkedPromise = r.parked.then( + () => 'resolved', + (error: Error) => `rejected:${error.message}`, + ); + + const expired = engine.expireRequest('t1', r.event.requestId, 'permission timed out'); + + assert.deepEqual(expired, { category: 'file_write', toolUseId: 'tu1' }); + expect(engine.pendingCount('t1')).toBe(0); + expect(await parkedPromise).toBe('rejected:permission timed out'); + expect(engine.recordResponse('t1', { requestId: r.event.requestId, decision: 'allow' })).toBeNull(); + }); + test('beginTurn is idempotent', () => { const { engine } = makeEngine(); engine.beginTurn('t1'); diff --git a/packages/runtime/src/__tests__/stream-watchdog.test.ts b/packages/runtime/src/__tests__/stream-watchdog.test.ts index d63df6a97f..34cac16327 100644 --- a/packages/runtime/src/__tests__/stream-watchdog.test.ts +++ b/packages/runtime/src/__tests__/stream-watchdog.test.ts @@ -74,6 +74,36 @@ describe('StreamWatchdog', () => { expect(fired).toEqual([{ phase: 'idle', elapsedMs: 10_000 }]); }); + test('nested pauses require matching resumes before idle timeout restarts', () => { + const timers = fakeTimers(3_500); + const fired: StreamWatchdogTimeout[] = []; + const watchdog = new StreamWatchdog({ + now: timers.now, + setTimer: timers.setTimer, + clearTimer: timers.clearTimer, + connectTimeoutMs: 30_000, + idleTimeoutMs: 10_000, + onTimeout: (timeout) => fired.push(timeout), + }); + + watchdog.start(); + watchdog.markActivity(); + watchdog.pause(); + watchdog.pause(); + timers.advance(600_000); + expect(fired).toEqual([]); + + watchdog.resume(); + timers.advance(60_000); + expect(fired).toEqual([]); + + watchdog.resume(); + timers.advance(9_999); + expect(fired).toEqual([]); + timers.advance(1); + expect(fired).toEqual([{ phase: 'idle', elapsedMs: 10_000 }]); + }); + test('stop cancels the active timer', () => { const timers = fakeTimers(4_000); const fired: StreamWatchdogTimeout[] = []; diff --git a/packages/runtime/src/ai-sdk-backend.ts b/packages/runtime/src/ai-sdk-backend.ts index b9978be594..146ccd750e 100644 --- a/packages/runtime/src/ai-sdk-backend.ts +++ b/packages/runtime/src/ai-sdk-backend.ts @@ -152,6 +152,7 @@ export type ModelFactory = (input: ModelFactoryInput) => unknown; export const TOOL_ERROR_RESULT_MAX_CHARS = 4000; export const INVALID_TOOL_NAME = 'invalid'; export const MAX_ACTIVE_SUBAGENT_TOOLS_PER_TURN = 5; +export const DEFAULT_PERMISSION_TIMEOUT_MS = 300_000; const SUBAGENT_TOOL_LIMIT_MESSAGE = '只读探索并发过多:同一轮最多 5 个子代理。请等待已有探索完成后再继续。'; export interface RepairableAiSdkToolCall { @@ -204,6 +205,8 @@ export interface AiSdkBackendInput { streamConnectTimeoutMs?: number; /** Timeout between SDK/tool events; paused while waiting on permission. Default 120s. */ streamIdleTimeoutMs?: number; + /** Timeout for a renderer/user permission decision. Default 300s. */ + permissionTimeoutMs?: number; /** Optional system prompt (skills + workspace AGENTS.md merged upstream). */ systemPrompt?: string | ((context: SystemPromptContext) => string | undefined | Promise); /** Provider-native options passed through to ai-sdk. */ @@ -275,7 +278,7 @@ export class AiSdkBackend implements AgentBackend { let thinkingText = ''; let thinkingSignature: string | undefined; const startedAt = this.now(); - let tokenUsage: { promptTokens?: number; completionTokens?: number; totalTokens?: number } | undefined; + let tokenUsage: NormalizedAiSdkUsage | undefined; let streamStatus: LlmCallRecord['status'] = 'success'; let streamErrorClass: string | undefined; @@ -434,17 +437,17 @@ export class AiSdkBackend implements AgentBackend { // Final usage event (await result.usage which resolves once stream ends). try { - const usage = await result.usage; - tokenUsage = usage; - if (usage) { + tokenUsage = normalizeAiSdkUsage(await result.usage); + if (tokenUsage) { const tu: TokenUsageMessage = { type: 'token_usage', id: this.newId(), turnId, ts: this.now(), - input: usage.promptTokens ?? 0, - output: usage.completionTokens ?? 0, - ...(usage.totalTokens !== undefined ? {} : {}), + input: tokenUsage.inputTokens, + output: tokenUsage.outputTokens, + ...(tokenUsage.cachedInputTokens > 0 ? { cacheRead: tokenUsage.cachedInputTokens } : {}), + ...(tokenUsage.cacheWriteInputTokens > 0 ? { cacheCreation: tokenUsage.cacheWriteInputTokens } : {}), }; await this.input.appendMessage(tu).catch(() => {}); queue.push({ @@ -452,8 +455,10 @@ export class AiSdkBackend implements AgentBackend { id: this.newId(), turnId, ts: this.now(), - input: usage.promptTokens ?? 0, - output: usage.completionTokens ?? 0, + input: tokenUsage.inputTokens, + output: tokenUsage.outputTokens, + ...(tokenUsage.cachedInputTokens > 0 ? { cacheRead: tokenUsage.cachedInputTokens } : {}), + ...(tokenUsage.cacheWriteInputTokens > 0 ? { cacheCreation: tokenUsage.cacheWriteInputTokens } : {}), } satisfies TokenUsageEvent); } } catch { @@ -507,8 +512,11 @@ export class AiSdkBackend implements AgentBackend { connectionSlug: this.input.connection.slug, providerId: this.input.connection.providerType, modelId: this.input.modelId, - inputTokens: tokenUsage?.promptTokens ?? 0, - outputTokens: tokenUsage?.completionTokens ?? 0, + inputTokens: tokenUsage?.inputTokens ?? 0, + outputTokens: tokenUsage?.outputTokens ?? 0, + cachedInputTokens: tokenUsage?.cachedInputTokens ?? 0, + cacheWriteInputTokens: tokenUsage?.cacheWriteInputTokens ?? 0, + reasoningTokens: tokenUsage?.reasoningTokens ?? 0, totalTokens: tokenUsage?.totalTokens, latencyMs: Math.max(0, this.now() - startedAt), status: streamStatus, @@ -595,11 +603,8 @@ export class AiSdkBackend implements AgentBackend { queue.push(verdict.event); let response: PermissionDecision; try { - this.currentWatchdog?.pause(); - response = await verdict.parked; - this.currentWatchdog?.resume(); + response = await this.awaitPermissionDecision(verdict, turnId); } catch (err) { - this.currentWatchdog?.resume(); const msg = formatSyntheticToolErrorText(err); const reason = formatSyntheticToolErrorText(`Permission flow aborted: ${msg}`); await this.writeSyntheticToolResult(toolUseId, turnId, reason, queue); @@ -1040,6 +1045,32 @@ export class AiSdkBackend implements AgentBackend { for await (const ev of queue) yield ev; } + private async awaitPermissionDecision( + verdict: Extract, { kind: 'prompt' }>, + turnId: string, + ): Promise { + const timeoutMs = this.input.permissionTimeoutMs ?? DEFAULT_PERMISSION_TIMEOUT_MS; + this.currentWatchdog?.pause(); + try { + if (timeoutMs <= 0) return await verdict.parked; + let timer: ReturnType | undefined; + const timeout = new Promise((_resolve, reject) => { + timer = setTimeout(() => { + const reason = `Permission request ${verdict.event.requestId} timed out after ${timeoutMs}ms`; + this.input.permissionEngine.expireRequest(turnId, verdict.event.requestId, reason); + reject(new Error(reason)); + }, timeoutMs); + }); + try { + return await Promise.race([verdict.parked, timeout]); + } finally { + if (timer !== undefined) clearTimeout(timer); + } + } finally { + this.currentWatchdog?.resume(); + } + } + private cleanupAfterTurn(turnId: string): void { this.input.permissionEngine.endTurn(turnId, this.aborted ? 'aborted' : 'completed'); this.abortController = null; @@ -1063,17 +1094,83 @@ interface AiSdkStreamChunk { toolName?: string; args?: unknown; result?: unknown; - usage?: { promptTokens?: number; completionTokens?: number; totalTokens?: number }; + usage?: AiSdkUsageLike; finishReason?: string; error?: unknown; } interface StreamTextResult { fullStream: AsyncIterable; - usage: Promise<{ promptTokens?: number; completionTokens?: number; totalTokens?: number } | undefined>; + usage: Promise; finishReason: Promise; } +interface AiSdkUsageLike { + promptTokens?: number; + completionTokens?: number; + totalTokens?: number; + inputTokens?: number; + outputTokens?: number; + cachedInputTokens?: number; + cacheWriteInputTokens?: number; + reasoningTokens?: number; + cacheReadInputTokens?: number; + cacheCreationInputTokens?: number; + inputTokenDetails?: { + cachedTokens?: number; + cacheReadTokens?: number; + cacheWriteTokens?: number; + reasoningTokens?: number; + }; + outputTokenDetails?: { + reasoningTokens?: number; + }; +} + +interface NormalizedAiSdkUsage { + inputTokens: number; + outputTokens: number; + cachedInputTokens: number; + cacheWriteInputTokens: number; + reasoningTokens: number; + totalTokens: number; +} + +export function normalizeAiSdkUsage(usage: AiSdkUsageLike | undefined): NormalizedAiSdkUsage | undefined { + if (!usage) return undefined; + const inputTokens = finiteToken(usage.inputTokens) ?? finiteToken(usage.promptTokens) ?? 0; + const outputTokens = finiteToken(usage.outputTokens) ?? finiteToken(usage.completionTokens) ?? 0; + const cachedInputTokens = + finiteToken(usage.cachedInputTokens) + ?? finiteToken(usage.cacheReadInputTokens) + ?? finiteToken(usage.inputTokenDetails?.cacheReadTokens) + ?? finiteToken(usage.inputTokenDetails?.cachedTokens) + ?? 0; + const cacheWriteInputTokens = + finiteToken(usage.cacheWriteInputTokens) + ?? finiteToken(usage.cacheCreationInputTokens) + ?? finiteToken(usage.inputTokenDetails?.cacheWriteTokens) + ?? 0; + const reasoningTokens = + finiteToken(usage.reasoningTokens) + ?? finiteToken(usage.outputTokenDetails?.reasoningTokens) + ?? finiteToken(usage.inputTokenDetails?.reasoningTokens) + ?? 0; + const totalTokens = finiteToken(usage.totalTokens) ?? inputTokens + outputTokens; + return { + inputTokens, + outputTokens, + cachedInputTokens, + cacheWriteInputTokens, + reasoningTokens, + totalTokens, + }; +} + +function finiteToken(value: unknown): number | undefined { + return typeof value === 'number' && Number.isFinite(value) && value >= 0 ? value : undefined; +} + function classifyError(error: unknown): string { if (!(error instanceof Error)) return 'Other'; const code = 'code' in error ? String((error as { code?: unknown }).code) : ''; @@ -1114,7 +1211,9 @@ function deriveToolResultStatus(content: ToolResultContent): ToolInvocationRecor } if (content.kind === 'rive_workflow' && content.ok === false) return 'error'; if (content.kind === 'web_search_error') return 'error'; - if (content.kind === 'office_document' && content.ok === false) return 'error'; + if (content.kind === 'office_document' && content.ok === false) { + return content.reason === 'officecli_aborted' ? 'aborted' : 'error'; + } return 'success'; } diff --git a/packages/runtime/src/permission-engine.ts b/packages/runtime/src/permission-engine.ts index 8371215bd5..1b8ac429c7 100644 --- a/packages/runtime/src/permission-engine.ts +++ b/packages/runtime/src/permission-engine.ts @@ -221,6 +221,21 @@ export class PermissionEngine { return { category: parked.category, toolUseId: parked.toolUseId }; } + /** + * Fail one parked request without ending the whole turn. + * Used by runtime-level permission timeouts so late UI responses do not + * resolve a tool call that has already failed closed. + */ + expireRequest(turnId: string, requestId: string, reason: string): { category: ToolCategory; toolUseId: string } | null { + const state = this.turns.get(turnId); + if (!state) return null; + const parked = state.parked.get(requestId); + if (!parked) return null; + state.parked.delete(requestId); + parked.reject(new Error(reason)); + return { category: parked.category, toolUseId: parked.toolUseId }; + } + /** Test/debug accessor. */ pendingCount(turnId: string): number { return this.turns.get(turnId)?.parked.size ?? 0; diff --git a/packages/runtime/src/stream-watchdog.ts b/packages/runtime/src/stream-watchdog.ts index 25784d907e..16f8353a4c 100644 --- a/packages/runtime/src/stream-watchdog.ts +++ b/packages/runtime/src/stream-watchdog.ts @@ -38,7 +38,7 @@ export class StreamWatchdog { private startedAt = 0; private lastActivityAt = 0; private sawActivity = false; - private paused = false; + private pauseCount = 0; private stopped = false; private timer: unknown; @@ -57,7 +57,7 @@ export class StreamWatchdog { this.startedAt = now; this.lastActivityAt = now; this.sawActivity = false; - this.paused = false; + this.pauseCount = 0; this.schedule(this.connectTimeoutMs); } @@ -65,18 +65,19 @@ export class StreamWatchdog { if (this.stopped) return; this.sawActivity = true; this.lastActivityAt = this.now(); - if (!this.paused) this.schedule(this.idleTimeoutMs); + if (this.pauseCount === 0) this.schedule(this.idleTimeoutMs); } pause(): void { if (this.stopped) return; - this.paused = true; + this.pauseCount += 1; this.clear(); } resume(): void { if (this.stopped) return; - this.paused = false; + this.pauseCount = Math.max(0, this.pauseCount - 1); + if (this.pauseCount > 0) return; this.markActivity(); } @@ -99,7 +100,7 @@ export class StreamWatchdog { } private fire(): void { - if (this.stopped || this.paused) return; + if (this.stopped || this.pauseCount > 0) return; this.stopped = true; this.clear(); const phase: StreamWatchdogPhase = this.sawActivity ? 'idle' : 'connect'; diff --git a/packages/storage/src/__tests__/session-store.test.ts b/packages/storage/src/__tests__/session-store.test.ts index 1878c9bcf9..fd5aab6c0a 100644 --- a/packages/storage/src/__tests__/session-store.test.ts +++ b/packages/storage/src/__tests__/session-store.test.ts @@ -222,6 +222,94 @@ describe('FileSessionStore CRUD', () => { }); }); + test('recovers readable messages around a corrupt JSONL message line', async () => { + await withStore(async (store, workspaceRoot) => { + const sessionId = 'corrupt-middle-line'; + const sessionDir = join(workspaceRoot, 'sessions', sessionId); + await mkdir(sessionDir, { recursive: true }); + await writeFile( + join(sessionDir, 'session.jsonl'), + [ + JSON.stringify(makeRawHeader({ id: sessionId, workspaceRoot, name: 'Corrupt middle' })), + JSON.stringify({ type: 'user', id: 'u1', turnId: 't1', ts: 2, text: 'hello' }), + '{"type":"assistant","id":"broken"', + JSON.stringify({ type: 'assistant', id: 'a1', turnId: 't1', ts: 4, text: 'recovered answer', modelId: 'fake' }), + '', + ].join('\n'), + 'utf8', + ); + + const messages = await store.readMessages(sessionId); + assert.equal(messages.length, 3); + assert.equal(messages[0]?.type, 'user'); + const note = messages[1]; + assert.equal(note?.type, 'system_note'); + if (note?.type !== 'system_note') throw new Error('corruption note missing'); + assert.equal(note.kind, 'error'); + assert.equal((note.data as { code?: unknown }).code, 'jsonl_parse_error'); + assert.equal((note.data as { lineNumber?: unknown }).lineNumber, 3); + assert.equal(typeof (note.data as { message?: unknown }).message, 'string'); + assert.ok(((note.data as { message?: string }).message ?? '').length > 0); + assert.equal(messages[2]?.type, 'assistant'); + + const [summary] = await store.list(); + assert.equal(summary?.id, sessionId); + assert.equal(summary?.lastMessagePreview, 'recovered answer'); + }); + }); + + test('silently drops a truncated tail JSONL message line', async () => { + await withStore(async (store, workspaceRoot) => { + const sessionId = 'truncated-tail-line'; + const sessionDir = join(workspaceRoot, 'sessions', sessionId); + await mkdir(sessionDir, { recursive: true }); + await writeFile( + join(sessionDir, 'session.jsonl'), + [ + JSON.stringify(makeRawHeader({ id: sessionId, workspaceRoot, name: 'Truncated tail' })), + JSON.stringify({ type: 'user', id: 'u1', turnId: 't1', ts: 2, text: 'survives' }), + '{"type":"assistant","id":"partial"', + ].join('\n'), + 'utf8', + ); + + const messages = await store.readMessages(sessionId); + assert.deepEqual(messages.map((message) => message.type), ['user']); + + const [summary] = await store.list(); + assert.equal(summary?.id, sessionId); + assert.equal(summary?.lastMessagePreview, 'survives'); + }); + }); + + test('reports a corrupt tail JSONL message line when it was newline-terminated', async () => { + await withStore(async (store, workspaceRoot) => { + const sessionId = 'corrupt-terminated-tail-line'; + const sessionDir = join(workspaceRoot, 'sessions', sessionId); + await mkdir(sessionDir, { recursive: true }); + await writeFile( + join(sessionDir, 'session.jsonl'), + [ + JSON.stringify(makeRawHeader({ id: sessionId, workspaceRoot, name: 'Corrupt terminated tail' })), + JSON.stringify({ type: 'user', id: 'u1', turnId: 't1', ts: 2, text: 'survives' }), + '{"type":"assistant","id":"durably-broken"', + '', + ].join('\n'), + 'utf8', + ); + + const messages = await store.readMessages(sessionId); + assert.equal(messages.length, 2); + assert.equal(messages[0]?.type, 'user'); + const note = messages[1]; + assert.equal(note?.type, 'system_note'); + if (note?.type !== 'system_note') throw new Error('corruption note missing'); + assert.equal(note.kind, 'error'); + assert.equal((note.data as { code?: unknown }).code, 'jsonl_parse_error'); + assert.equal((note.data as { lineNumber?: unknown }).lineNumber, 3); + }); + }); + test('derives lastMessagePreview from visible user and assistant messages', async () => { await withStore(async (store) => { const header = await store.create(makeInput({ name: 'Preview' })); @@ -460,6 +548,30 @@ function makeInput(overrides: Partial = {}): CreateSessionIn }; } +function makeRawHeader(overrides: Partial = {}): SessionHeader { + return { + id: 'raw-session', + workspaceRoot: '/tmp/workspace', + cwd: '/tmp/cwd', + createdAt: 1, + lastUsedAt: 1, + name: 'Raw session', + isFlagged: false, + labels: [], + isArchived: false, + status: 'active', + statusUpdatedAt: 1, + hasUnread: false, + backend: 'fake', + llmConnectionSlug: 'fake', + connectionLocked: false, + model: 'fake-model', + permissionMode: 'ask', + schemaVersion: 1, + ...overrides, + }; +} + async function withStore( fn: (store: ReturnType, workspaceRoot: string) => Promise, ): Promise { diff --git a/packages/storage/src/session-store.ts b/packages/storage/src/session-store.ts index 1978212da0..eb4df736b9 100644 --- a/packages/storage/src/session-store.ts +++ b/packages/storage/src/session-store.ts @@ -223,10 +223,23 @@ class FileSessionStore implements SessionStore { private async readFilePartsUnlocked(sessionId: string): Promise<{ header: SessionHeader; messages: StoredMessage[] }> { const text = await readFile(this.sessionPath(sessionId), 'utf8'); - const lines = text.split('\n').filter((line) => line.trim().length > 0); + const rawLines = text.split('\n'); + const endsWithNewline = text.endsWith('\n'); + const lines = rawLines + .map((line, index) => ({ line, lineNumber: index + 1 })) + .filter((entry) => entry.line.trim().length > 0); if (lines.length === 0 || !lines[0]) throw new Error(`Session ${sessionId} is empty`); - const header = migrateHeader(JSON.parse(lines[0]) as StoredSessionHeader); - const messages = lines.slice(1).map((line) => JSON.parse(line) as StoredMessage); + const header = migrateHeader(JSON.parse(lines[0].line) as StoredSessionHeader); + const messages: StoredMessage[] = []; + const lastLineNumber = lines.at(-1)?.lineNumber; + for (const entry of lines.slice(1)) { + try { + messages.push(JSON.parse(entry.line) as StoredMessage); + } catch (error) { + if (!endsWithNewline && entry.lineNumber === lastLineNumber) continue; + messages.push(createJsonlCorruptionNote(header, entry.lineNumber, error)); + } + } return { header, messages }; } @@ -269,6 +282,20 @@ type StoredSessionHeader = Omit 0 ? header.model : 'default';