From 92a4227de6f10d0247a37778122e1a5681bc8c54 Mon Sep 17 00:00:00 2001 From: Alex Alecu Date: Wed, 27 May 2026 17:17:41 +0300 Subject: [PATCH 1/4] docs(code-review): clarify durable object boundary --- services/cloud-agent-next/AGENTS.md | 6 +++ .../cloud-agent-next/src/session-service.ts | 3 +- services/cloud-agent/src/workspace.ts | 3 +- services/code-review-infra/AGENTS.md | 38 +++++++++++++++++++ 4 files changed, 48 insertions(+), 2 deletions(-) create mode 100644 services/code-review-infra/AGENTS.md diff --git a/services/cloud-agent-next/AGENTS.md b/services/cloud-agent-next/AGENTS.md index 4d78c45bf3..841186274e 100644 --- a/services/cloud-agent-next/AGENTS.md +++ b/services/cloud-agent-next/AGENTS.md @@ -129,6 +129,12 @@ This pattern blocks API endpoints from running for external contributors who don - Put Kilo/job behavior in `wrapper/` or Kilo SDK integration code when it does not require durable DO coordination. - Avoid growing `CloudAgentSession.ts` with product behavior that can live in the wrapper, Kilo SDK layer, or a small helper module. +### Code Review Integration Boundary + +- `services/code-review-infra/src/code-review-orchestrator.ts` owns code-review product and workflow semantics: review IDs and attempt state, status reconciliation, retry or fresh-session selection, review callback construction, cancellation policy, and review-specific dispatch rules. +- `CloudAgentSession` may provide reusable durable runtime operations used by code review, including the retained `prepareSession` / `updateSession` / `sendMessageV2` callback seam and runtime enforcement of the read-only code-review command guard. It must not acquire review IDs or attempt state, review status reconciliation, retry or fresh-session selection, review callback construction, or review-specific dispatch rules. +- When code-review integration needs another generic session or runtime capability, expose a narrow reusable primitive here and leave the code-review decision and policy in `CodeReviewOrchestrator`. + ### Runtime Guidelines - Durable Object calls should be retried using `withDORetry` in `src/utils/do-retry.ts` diff --git a/services/cloud-agent-next/src/session-service.ts b/services/cloud-agent-next/src/session-service.ts index 15c8b42b3e..870fcc02fa 100644 --- a/services/cloud-agent-next/src/session-service.ts +++ b/services/cloud-agent-next/src/session-service.ts @@ -73,7 +73,8 @@ import { normalizeAgentMode } from './schema.js'; const SETUP_COMMAND_TIMEOUT_SECONDS = 300; // 5 minutes const DEFAULT_DENIED_COMMAND_PATTERNS = ['rm -rf', 'sudo rm', 'mkfs', 'dd if=']; -// Keep in sync with: cloudflare-code-review-infra/src/code-review-orchestrator.ts +// This service enforces runtime permissions for code-review sessions. +// `services/code-review-infra/src/code-review-orchestrator.ts` maintains only the review-side observability subset. // mkdir and touch are intentionally allowed for agent scratch space during analysis const CODE_REVIEW_ALLOWED_COMMANDS = [ 'ls', diff --git a/services/cloud-agent/src/workspace.ts b/services/cloud-agent/src/workspace.ts index 7611cf156c..fb57e33e5b 100644 --- a/services/cloud-agent/src/workspace.ts +++ b/services/cloud-agent/src/workspace.ts @@ -59,7 +59,8 @@ const DEFAULT_ALLOWED_COMMANDS = [ const DEFAULT_DENIED_COMMAND_PATTERNS = ['rm -rf', 'sudo rm', 'mkfs', 'dd if=']; -// Keep in sync with: cloud-agent-next/src/session-service.ts, cloudflare-code-review-infra/src/code-review-orchestrator.ts +// This service enforces legacy runtime permissions; keep them aligned with cloud-agent-next/src/session-service.ts. +// `services/code-review-infra/src/code-review-orchestrator.ts` maintains only the review-side observability subset. // mkdir and touch are intentionally allowed for agent scratch space during analysis const CODE_REVIEW_ALLOWED_COMMANDS = [ 'ls', diff --git a/services/code-review-infra/AGENTS.md b/services/code-review-infra/AGENTS.md new file mode 100644 index 0000000000..878868410d --- /dev/null +++ b/services/code-review-infra/AGENTS.md @@ -0,0 +1,38 @@ +# Code Review Infrastructure Guidance + +This guidance supplements the repository-level `AGENTS.md` rules for work in `services/code-review-infra/`. + +## Ownership Boundary + +`src/code-review-orchestrator.ts` contains the `CodeReviewOrchestrator` Durable Object. It owns execution orchestration after Next.js has dispatched a review. Next.js remains responsible for queue and concurrency business decisions before dispatch. + +Keep these code-review concerns in `CodeReviewOrchestrator`: + +- Review and attempt identity, lifecycle state, status synchronization, and the web status/callback contract. +- Selection and invocation of supported cloud-agent execution paths, including compatibility behavior for legacy SSE execution. +- Review-specific continuation, fresh-session retry or fallback, cancellation, terminal reconciliation, and other review-facing policy decisions. +- Construction and assignment of per-attempt callback targets when coordinating a dispatched review. + +Do not duplicate generic cloud-agent session/runtime mechanics in this service. In particular, do not add implementations of durable session state, message admission or queues, wrapper/socket supervision and fencing, sandbox execution primitives, session event replay, or callback delivery/outbox transport owned by `services/cloud-agent-next/`. + +## Cloud Agent Seam + +- Request narrowly scoped, reusable session/runtime capabilities through cloud-agent APIs and callback contracts; keep review-level decisions and review identity/status state in `CodeReviewOrchestrator`. +- The retained prepare/update/send flow may be used where the cloud-agent-next contract requires it for review continuation, but the review-specific reason for continuation or fresh retry belongs here. +- Cloud-agent services enforce execution permissions, including read-only code-review command policy. Any command-risk matching in this service is review-side observability or compatibility handling, not the canonical runtime guard. + +## Worker and Durable Object Safety + +- Follow the inherited Worker and Durable Object rules in the root `AGENTS.md`, including avoiding module-scope caching of transport-owning database or SDK clients and using approved per-use helpers. +- Preserve callback authentication and redaction boundaries; never log tokens, credentials, auth headers, cookies, or webhook secrets. +- Keep changes within this service's orchestration role unless a shared cloud-agent capability must be introduced through an explicit contract. + +## Targeted Validation + +From the repository root, use the narrowest applicable checks: + +- `pnpm --dir services/code-review-infra test` +- `pnpm --dir services/code-review-infra typecheck` +- `pnpm --dir services/code-review-infra lint` + +Run only the checks relevant to the files changed, and follow the root formatting requirement before committing. From 1cbed07512d26a8623d719f5d0030a3017b15ce2 Mon Sep 17 00:00:00 2001 From: Alex Alecu Date: Wed, 27 May 2026 21:41:03 +0300 Subject: [PATCH 2/4] feat(cloud-agent): enforce managed session policies --- packages/worker-utils/package.json | 3 +- .../src/cloud-agent-next-client.ts | 53 ++++ packages/worker-utils/src/index.ts | 11 + .../src/managed-session-policy.ts | 237 ++++++++++++++++++ .../src/balance-validation.ts | 1 + .../cloud-agent-next/src/callbacks/types.ts | 13 + .../src/execution/orchestrator.ts | 17 +- .../cloud-agent-next/src/execution/types.ts | 7 + .../src/middleware/balance.ts | 1 + .../src/persistence/CloudAgentSession.ts | 15 +- .../src/persistence/execution-policy.ts | 62 +++++ .../src/persistence/session-metadata.ts | 2 + .../src/router/handlers/session-execution.ts | 160 +++++++----- .../src/router/handlers/session-management.ts | 3 + .../src/router/handlers/session-prepare.ts | 1 + .../src/router/handlers/session-questions.ts | 36 +++ .../cloud-agent-next/src/router/schemas.ts | 51 ++++ .../src/session-service.test.ts | 37 +++ .../cloud-agent-next/src/session-service.ts | 61 ++--- .../src/session/message-settlement-outbox.ts | 38 ++- .../src/session/pending-messages.ts | 13 +- .../src/session/queue-message.ts | 1 + .../src/session/session-message-queue.test.ts | 26 ++ .../src/session/session-message-queue.ts | 94 ++++++- .../src/session/session-message-state.ts | 85 ++++++- .../src/session/session-registration.ts | 1 + .../src/session/session-requests.ts | 4 + .../src/shared/managed-session-policy.ts | 88 +++++++ .../src/shared/wrapper-bootstrap.ts | 3 + .../src/terminal/access.test.ts | 21 +- .../cloud-agent-next/src/terminal/access.ts | 15 ++ .../test/unit/wrapper/server.test.ts | 70 ++++++ .../wrapper/src/connection.ts | 74 ++++-- .../cloud-agent-next/wrapper/src/server.ts | 51 ++++ .../cloud-agent-next/wrapper/src/state.ts | 10 + 35 files changed, 1212 insertions(+), 153 deletions(-) create mode 100644 packages/worker-utils/src/managed-session-policy.ts create mode 100644 services/cloud-agent-next/src/persistence/execution-policy.ts create mode 100644 services/cloud-agent-next/src/shared/managed-session-policy.ts diff --git a/packages/worker-utils/package.json b/packages/worker-utils/package.json index 0f8df984fe..ad2539ccc4 100644 --- a/packages/worker-utils/package.json +++ b/packages/worker-utils/package.json @@ -15,7 +15,8 @@ "./kilo-pass-bonus-projection": "./src/kilo-pass-bonus-projection.ts", "./git-url": "./src/git-url.ts", "./callback-token": "./src/callback-token.ts", - "./kilo-model-id": "./src/kilo-model-id.ts" + "./kilo-model-id": "./src/kilo-model-id.ts", + "./managed-session-policy": "./src/managed-session-policy.ts" }, "scripts": { "test": "vitest run", diff --git a/packages/worker-utils/src/cloud-agent-next-client.ts b/packages/worker-utils/src/cloud-agent-next-client.ts index 2c0717bf4b..08022df837 100644 --- a/packages/worker-utils/src/cloud-agent-next-client.ts +++ b/packages/worker-utils/src/cloud-agent-next-client.ts @@ -17,6 +17,33 @@ export type CallbackTarget = { headers?: Record; }; +export type CloudAgentRuntimeAgent = { + slug: string; + name: string; + config: Record; +}; + +export type CloudAgentManagedSession = { + executionPolicy?: { + name: string; + permissionOverrides?: Record; + tools?: Record; + interaction?: { + question?: 'allow' | 'deny'; + permission?: 'allow' | 'auto-approve' | 'auto-reject'; + terminal?: 'allow' | 'deny'; + }; + runtime?: { + nonInteractive?: boolean; + }; + }; +}; + +export type CloudAgentMessageCompletion = { + callbackTarget?: CallbackTarget; + gateThreshold?: 'off' | 'all' | 'warning' | 'critical'; +}; + export type CloudAgentPrepareSessionInput = { prompt: string; mode: string; @@ -30,10 +57,12 @@ export type CloudAgentPrepareSessionInput = { kilocodeOrganizationId?: string; envVars?: Record; mcpServers?: Record; + runtimeAgents?: CloudAgentRuntimeAgent[]; upstreamBranch?: string; callbackTarget?: CallbackTarget; createdOnPlatform?: string; gateThreshold?: 'off' | 'all' | 'warning' | 'critical'; + managedSession?: CloudAgentManagedSession; }; export type CloudAgentPrepareSessionOutput = { @@ -66,6 +95,10 @@ export type CloudAgentSendMessageInput = { gitToken?: string; }; +export type CloudAgentSendMessageInternalInput = CloudAgentSendMessageInput & { + completion?: CloudAgentMessageCompletion; +}; + export type CloudAgentSendMessageOutput = { executionId: string; status?: string; @@ -239,6 +272,11 @@ export type CloudAgentNextFetchClient = { input: CloudAgentSendMessageInput ): Promise; + sendMessageV2Internal( + headers: Record, + input: CloudAgentSendMessageInternalInput + ): Promise; + getSessionHealth( headers: Record, input: CloudAgentSessionHealthInput @@ -310,6 +348,21 @@ export function createCloudAgentNextFetchClient(baseUrl: string): CloudAgentNext return data as unknown as CloudAgentSendMessageOutput; }, + async sendMessageV2Internal(headers, input) { + const data = await trpcPost>( + trpc('sendMessageV2Internal'), + headers, + input, + 'sendMessageV2Internal' + ); + if (typeof data.executionId !== 'string') { + throw new Error( + `Unexpected sendMessageV2Internal response shape: ${JSON.stringify(data).slice(0, 500)}` + ); + } + return data as unknown as CloudAgentSendMessageOutput; + }, + async getSessionHealth(headers, input) { const data = await trpcPost>( trpc('getSessionHealth'), diff --git a/packages/worker-utils/src/index.ts b/packages/worker-utils/src/index.ts index 2f37f71264..0342900805 100644 --- a/packages/worker-utils/src/index.ts +++ b/packages/worker-utils/src/index.ts @@ -47,6 +47,17 @@ export type { } from './cloud-agent-next-client.js'; export { CloudAgentNextBillingError, CloudAgentNextError } from './cloud-agent-next-client.js'; +export { + CODE_REVIEW_RUNTIME_AGENT_SLUG, + buildCodeReviewManagedSessionPolicy, + buildCodeReviewRuntimeAgent, +} from './managed-session-policy.js'; +export type { + CloudAgentRuntimeAgent, + ManagedSessionExecutionPolicy, + ManagedSessionInteractionPolicy, +} from './managed-session-policy.js'; + export { signKiloToken, verifyKiloToken, diff --git a/packages/worker-utils/src/managed-session-policy.ts b/packages/worker-utils/src/managed-session-policy.ts new file mode 100644 index 0000000000..0ca51c8aea --- /dev/null +++ b/packages/worker-utils/src/managed-session-policy.ts @@ -0,0 +1,237 @@ +export type ManagedSessionInteractionPolicy = { + question?: 'allow' | 'deny'; + permission?: 'allow' | 'auto-approve' | 'auto-reject'; + terminal?: 'allow' | 'deny'; +}; + +export type ManagedSessionExecutionPolicy = { + name: string; + permissionOverrides?: Record; + tools?: Record; + interaction?: ManagedSessionInteractionPolicy; + runtime?: { + nonInteractive?: boolean; + }; +}; + +export type CloudAgentRuntimeAgent = { + slug: string; + name: string; + config: { + prompt?: string; + description?: string; + mode?: 'subagent' | 'primary' | 'all'; + model?: string | null; + variant?: string; + permission?: 'allow' | 'ask' | 'deny' | Record; + options?: Record; + }; +}; + +export const CODE_REVIEW_RUNTIME_AGENT_SLUG = 'code-review'; + +const DEFAULT_DENIED_COMMAND_PATTERNS = ['rm -rf', 'sudo rm', 'mkfs', 'dd if=']; + +// mkdir and touch are intentionally allowed for agent scratch space during analysis. +const CODE_REVIEW_ALLOWED_COMMANDS = [ + 'ls', + 'cat', + 'echo', + 'pwd', + 'find', + 'grep', + 'wc', + 'sort', + 'uniq', + 'cut', + 'tr', + 'nl', + 'jq', + 'git', + 'git fetch', + 'git pull', + 'gh pr diff', + 'gh pr view', + 'gh api repos/*/issues/*/comments --input*', + 'gh api repos/*/issues/comments/* -X PATCH*', + 'gh api repos/*/pulls/*/reviews --input*', + 'glab mr diff', + 'glab mr view', + 'glab api --method POST *merge_requests/*/notes*', + 'glab api --method PUT *merge_requests/*/notes/*', + 'glab api --method POST *merge_requests/*/discussions*', + 'whoami', + 'date', + 'stat', + 'file', + 'head', + 'tail', + 'sed', + 'cd', + 'mkdir', + 'touch', +]; + +const CODE_REVIEW_DENIED_COMMAND_PATTERNS = [ + 'bash', + 'sh', + 'zsh', + 'fish', + 'sed -i', + 'sed -*i', + 'sed --in-place', + 'sed --in-place*', + 'sed * -i', + 'sed * -*i', + 'sed * --in-place', + 'sed * --in-place*', + 'sort -o', + 'sort -o*', + 'sort -*o', + 'sort --output', + 'sort --output*', + 'sort * -o', + 'sort * -o*', + 'sort * -*o', + 'sort * --output', + 'sort * --output*', + 'uniq * *', + 'python', + 'python3', + 'node', + 'irb', + 'php -a', + 'rails console', + 'vi', + 'vim', + 'nvim', + 'nano', + 'emacs', + 'less', + 'more', + 'top', + 'htop', + 'watch', + 'tail -f', + 'ssh', + 'tmux', + 'screen', + 'git add', + 'git branch', + 'git clean', + 'git commit', + 'git config', + 'git mv', + 'git push', + 'git restore', + 'git rm', + 'git merge', + 'git rebase', + 'git cherry-pick', + 'git reset', + 'git checkout', + 'git switch', + 'git stash', + 'git tag', + 'git worktree', + 'git am', + 'git apply', + 'git remote set-url', + 'gh pr merge', + 'gh pr review', + 'gh pr create', + 'gh pr close', + 'gh pr edit', + 'gh pr checkout', + 'gh auth login', + 'gh auth refresh', + 'gh issue', + 'gh repo create', + 'gh repo fork', + 'glab auth', + 'glab mr approve', + 'glab mr close', + 'glab mr create', + 'glab mr delete', + 'glab mr merge', + 'glab mr reopen', + 'glab mr update', + 'glab repo', + 'glab issue', + 'glab pipeline', + 'glab release', + 'glab variable', + 'npm test', + 'pnpm test', + 'bun test', + 'yarn test', + 'pytest', + 'vitest', +]; + +function buildCommandGuardBashPermissions(params: { + allowed: readonly string[]; + denied: readonly string[]; +}): Record { + const bashPermissions: Record = {}; + for (const cmd of params.allowed) { + bashPermissions[cmd] = 'allow'; + bashPermissions[`${cmd} *`] = 'allow'; + } + for (const cmd of params.denied) { + bashPermissions[cmd] = 'deny'; + bashPermissions[`${cmd} *`] = 'deny'; + } + return bashPermissions; +} + +export function buildCodeReviewManagedSessionPolicy(): ManagedSessionExecutionPolicy { + return { + name: 'code-review-read-only', + runtime: { nonInteractive: true }, + interaction: { + question: 'deny', + permission: 'auto-reject', + terminal: 'deny', + }, + tools: { + question: false, + plan_enter: false, + plan_exit: false, + }, + permissionOverrides: { + read: 'allow', + edit: 'deny', + bash: buildCommandGuardBashPermissions({ + allowed: CODE_REVIEW_ALLOWED_COMMANDS, + denied: [...DEFAULT_DENIED_COMMAND_PATTERNS, ...CODE_REVIEW_DENIED_COMMAND_PATTERNS], + }), + webfetch: 'deny', + websearch: 'deny', + codesearch: 'deny', + todowrite: 'allow', + todoread: 'allow', + question: 'deny', + }, + }; +} + +export function buildCodeReviewRuntimeAgent(options: { + model: string; + variant?: string; + permission?: Record; +}): CloudAgentRuntimeAgent { + return { + slug: CODE_REVIEW_RUNTIME_AGENT_SLUG, + name: 'Code Review', + config: { + mode: 'primary', + model: options.model, + variant: options.variant, + description: 'Read-only, non-interactive code review agent.', + prompt: + 'You are Kilo Code performing a read-only pull request review. Inspect the repository and diff, do not modify files, and follow the user prompt for review output.', + ...(options.permission ? { permission: options.permission } : {}), + }, + }; +} diff --git a/services/cloud-agent-next/src/balance-validation.ts b/services/cloud-agent-next/src/balance-validation.ts index 26f913de94..ac5833f468 100644 --- a/services/cloud-agent-next/src/balance-validation.ts +++ b/services/cloud-agent-next/src/balance-validation.ts @@ -139,6 +139,7 @@ export async function fetchOrgIdForSession( export const BALANCE_REQUIRED_MUTATIONS = new Set([ 'initiateFromKilocodeSessionV2', 'sendMessageV2', + 'sendMessageV2Internal', 'start', 'send', ]); diff --git a/services/cloud-agent-next/src/callbacks/types.ts b/services/cloud-agent-next/src/callbacks/types.ts index 5172270fbb..8f3da44b46 100644 --- a/services/cloud-agent-next/src/callbacks/types.ts +++ b/services/cloud-agent-next/src/callbacks/types.ts @@ -21,6 +21,19 @@ export type ExecutionCallbackPayload = { * Undefined when no assistant message has been recorded yet. */ lastAssistantMessageText?: string; + completionData?: { + assistantMessageId?: string; + lastAssistantMessageText?: string; + gateResult?: 'pass' | 'fail'; + lastSeenBranch?: string; + completedAt?: number; + source?: + | 'assistant_message_event' + | 'idle_reconciliation' + | 'wrapper_failure' + | 'interrupt' + | 'delivery_failure'; + }; /** * Deterministic idempotency key based on messageId. * Receivers can use this to safely deduplicate retried callbacks after a diff --git a/services/cloud-agent-next/src/execution/orchestrator.ts b/services/cloud-agent-next/src/execution/orchestrator.ts index c071fa7d42..650984a531 100644 --- a/services/cloud-agent-next/src/execution/orchestrator.ts +++ b/services/cloud-agent-next/src/execution/orchestrator.ts @@ -34,16 +34,11 @@ import { withPreparationInfrastructureRecovery, } from '../sandbox-recovery.js'; import { checkDiskAndCleanBeforeSetup } from '../workspace.js'; +import { resolveSessionExecutionPolicy } from '../persistence/execution-policy.js'; /** Maximum time allowed for workspace preparation (resume, init, fast path). */ const PREPARE_WORKSPACE_TIMEOUT_MS = 10 * 60 * 1000; -const CODE_REVIEW_DISABLED_TOOLS = { - question: false, - plan_enter: false, - plan_exit: false, -} satisfies Record; - function withWorkspacePreparationTimeout(operation: Promise, step: string): Promise { return withTimeout( operation, @@ -400,17 +395,9 @@ export class ExecutionOrchestrator { } } - private getCreatedOnPlatform( - plan: FencedWrapperDispatchRequest | FencedLegacyExecutionRequest - ): string | undefined { - return plan.workspace.metadata?.identity?.createdOnPlatform; - } - private getToolOverrides( plan: FencedWrapperDispatchRequest | FencedLegacyExecutionRequest ): Record | undefined { - return this.getCreatedOnPlatform(plan) === 'code-review' - ? CODE_REVIEW_DISABLED_TOOLS - : undefined; + return resolveSessionExecutionPolicy(plan.workspace.metadata)?.tools; } } diff --git a/services/cloud-agent-next/src/execution/types.ts b/services/cloud-agent-next/src/execution/types.ts index e02e77bce7..7ba74f0708 100644 --- a/services/cloud-agent-next/src/execution/types.ts +++ b/services/cloud-agent-next/src/execution/types.ts @@ -11,6 +11,7 @@ import type { AgentMode } from '../schema.js'; import type { Images } from '../router/schemas.js'; import type { SessionMetadata } from '../persistence/session-metadata.js'; import type { CloudAgentSessionState } from '../persistence/types.js'; +import type { CallbackTarget } from '../callbacks/index.js'; // --------------------------------------------------------------------------- // Execution Modes @@ -105,6 +106,11 @@ export type TurnFinalization = { condenseOnComplete?: boolean; }; +export type MessageCompletionPolicy = { + callbackTarget?: CallbackTarget; + gateThreshold?: 'off' | 'all' | 'warning' | 'critical'; +}; + /** Session policy extends per-turn finalization with session-only gates. */ export type SessionFinalization = TurnFinalization & { gateThreshold?: 'off' | 'all' | 'warning' | 'critical'; @@ -194,6 +200,7 @@ export type QueueExecutionTurnCommand = { turn: ExecutionTurnSubmission; agent?: AgentSelectionOverride; finalization?: TurnFinalization; + completion?: MessageCompletionPolicy; }; /** Current-path submitted message before durable admission resolves identity/defaults. */ diff --git a/services/cloud-agent-next/src/middleware/balance.ts b/services/cloud-agent-next/src/middleware/balance.ts index 8f7d67f1e4..ee6f54027e 100644 --- a/services/cloud-agent-next/src/middleware/balance.ts +++ b/services/cloud-agent-next/src/middleware/balance.ts @@ -68,6 +68,7 @@ export const balanceMiddleware = createMiddleware( // carries `kilocodeOrganizationId` in the body, so it is not included here. if ( (procedureName === 'sendMessageV2' || + procedureName === 'sendMessageV2Internal' || procedureName === 'initiateFromKilocodeSessionV2' || procedureName === 'send') && !orgId && diff --git a/services/cloud-agent-next/src/persistence/CloudAgentSession.ts b/services/cloud-agent-next/src/persistence/CloudAgentSession.ts index 25809d206f..bd5fd74331 100644 --- a/services/cloud-agent-next/src/persistence/CloudAgentSession.ts +++ b/services/cloud-agent-next/src/persistence/CloudAgentSession.ts @@ -196,6 +196,7 @@ type GroupedRegisterSessionInput = { }; profile?: SessionProfileBundle; finalization?: SessionFinalization; + executionPolicy?: SessionMetadata['executionPolicy']; callback?: SessionMetadata['callback']; workspace?: Pick< NonNullable, @@ -238,7 +239,8 @@ function isSameInitialAdmissionConfiguration( metadata.agent.model === input.agent.model && metadata.agent.variant === input.agent.variant && metadata.finalization?.autoCommit === input.finalization?.autoCommit && - metadata.finalization?.condenseOnComplete === input.finalization?.condenseOnComplete + metadata.finalization?.condenseOnComplete === input.finalization?.condenseOnComplete && + JSON.stringify(metadata.executionPolicy) === JSON.stringify(input.executionPolicy) ); } @@ -306,6 +308,16 @@ export class CloudAgentSession extends DurableObject { kiloSessionId: metadata.auth.kiloSessionId, gateResult, lastAssistantMessageText, + completionData: + status === 'completed' + ? { + ...(gateResult !== undefined ? { gateResult } : {}), + ...(lastAssistantMessageText !== undefined ? { lastAssistantMessageText } : {}), + ...(metadata.repository?.upstreamBranch !== undefined + ? { lastSeenBranch: metadata.repository.upstreamBranch } + : {}), + } + : undefined, }; if (messageId) { @@ -1386,6 +1398,7 @@ export class CloudAgentSession extends DurableObject { appendSystemPrompt: input.agent.appendSystemPrompt, }, finalization: input.finalization, + executionPolicy: input.executionPolicy, profile: input.profile, callback: input.callback, workspace: input.workspace, diff --git a/services/cloud-agent-next/src/persistence/execution-policy.ts b/services/cloud-agent-next/src/persistence/execution-policy.ts new file mode 100644 index 0000000000..8b9b34ea10 --- /dev/null +++ b/services/cloud-agent-next/src/persistence/execution-policy.ts @@ -0,0 +1,62 @@ +import { PermissionConfigSchema } from '@kilocode/db/schema-types'; +import * as z from 'zod'; +import { buildCodeReviewManagedSessionPolicy } from '@kilocode/worker-utils/managed-session-policy'; + +const PermissionOverridesSchema = z.intersection( + PermissionConfigSchema, + z.record(z.string(), z.unknown()) +); + +export const ManagedSessionExecutionPolicySchema = z + .object({ + name: z + .string() + .min(1) + .max(100) + .regex(/^[a-z][a-z0-9-]*$/, 'Policy name must be a slug'), + permissionOverrides: PermissionOverridesSchema.optional(), + tools: z.record(z.string().min(1).max(100), z.boolean()).optional(), + interaction: z + .object({ + question: z.enum(['allow', 'deny']).optional(), + permission: z.enum(['allow', 'auto-approve', 'auto-reject']).optional(), + terminal: z.enum(['allow', 'deny']).optional(), + }) + .strict() + .optional(), + runtime: z + .object({ + nonInteractive: z.boolean().optional(), + }) + .strict() + .optional(), + }) + .strict(); + +export type ManagedSessionExecutionPolicy = z.infer; + +type SessionExecutionPolicySource = { + executionPolicy?: ManagedSessionExecutionPolicy; + identity?: { + createdOnPlatform?: string; + }; +}; + +const LegacyCodeReviewExecutionPolicy = ManagedSessionExecutionPolicySchema.parse( + buildCodeReviewManagedSessionPolicy() +); + +export function resolveSessionExecutionPolicy( + metadata: SessionExecutionPolicySource | null | undefined +): ManagedSessionExecutionPolicy | undefined { + if (!metadata) return undefined; + if (metadata.executionPolicy) return metadata.executionPolicy; + + // Legacy code-review sessions predate the explicit managed-session policy. + // Keep them constrained until their Durable Object state naturally expires. + if (metadata.identity?.createdOnPlatform === 'code-review') { + return LegacyCodeReviewExecutionPolicy; + } + + return undefined; +} diff --git a/services/cloud-agent-next/src/persistence/session-metadata.ts b/services/cloud-agent-next/src/persistence/session-metadata.ts index ea13bcd1cc..7cfbd002eb 100644 --- a/services/cloud-agent-next/src/persistence/session-metadata.ts +++ b/services/cloud-agent-next/src/persistence/session-metadata.ts @@ -9,6 +9,7 @@ import { MetadataSchema as LegacySessionMetadataSchema, SessionProfileBundleSchema, } from './schemas.js'; +import { ManagedSessionExecutionPolicySchema } from './execution-policy.js'; const SandboxIdSchema = z .string() @@ -164,6 +165,7 @@ export const CurrentSessionMetadataSchema = z initialMessage: MetadataInitialMessageSchema.optional(), agent: MetadataAgentSchema.optional(), finalization: MetadataFinalizationSchema.optional(), + executionPolicy: ManagedSessionExecutionPolicySchema.optional(), profile: SessionProfileBundleSchema.optional(), callback: MetadataCallbackSchema.optional(), workspace: MetadataWorkspaceSchema.optional(), diff --git a/services/cloud-agent-next/src/router/handlers/session-execution.ts b/services/cloud-agent-next/src/router/handlers/session-execution.ts index 86ea01d7a0..56feecf6c4 100644 --- a/services/cloud-agent-next/src/router/handlers/session-execution.ts +++ b/services/cloud-agent-next/src/router/handlers/session-execution.ts @@ -12,14 +12,17 @@ * * New callers should prefer the unified `start` / `send` endpoints. */ -import { protectedProcedure } from '../auth.js'; +import type * as z from 'zod'; +import { internalApiProtectedProcedure, protectedProcedure } from '../auth.js'; import { logger, withLogTags } from '../../logger.js'; import { InitiateFromPreparedSessionInput, SendMessageV2Input, + TrustedSendMessageV2Input, LegacyExecutionResponse, } from '../schemas.js'; import type { SessionId } from '../../types/ids.js'; +import type { Env } from '../../types.js'; import { queueMessage, replayMessageIfAlreadyAdmitted } from '../../session/queue-message.js'; import { admitLegacyPreparedInitialMessage, @@ -36,6 +39,15 @@ import { preflightPreparedInitialPromptModel, } from '../../session/model-preflight.js'; +type SendMessageV2InputValue = z.infer; +type TrustedSendMessageV2InputValue = z.infer; + +type SendMessageV2Context = { + env: Env; + userId: string; + botId?: string; +}; + function withLegacyExecutionId(ack: QueueAckResponse): LegacyExecutionResponse { return { ...ack, @@ -43,6 +55,78 @@ function withLegacyExecutionId(ack: QueueAckResponse): LegacyExecutionResponse { }; } +async function queueSendMessageV2( + input: SendMessageV2InputValue | TrustedSendMessageV2InputValue, + ctx: SendMessageV2Context, + procedure: string +): Promise { + const commandPayload = + 'payload' in input && input.payload.type === 'command' ? input.payload : undefined; + const promptPayload = + 'prompt' in input + ? { + prompt: input.prompt, + mode: input.mode, + model: input.model, + variant: input.variant, + } + : 'payload' in input && input.payload.type === 'prompt' + ? input.payload + : undefined; + let turn: ExecutionTurnSubmission; + let agent: AgentSelectionOverride | undefined; + if (commandPayload) { + turn = { + type: 'command', + id: input.messageId ?? undefined, + command: commandPayload.command, + arguments: commandPayload.arguments, + images: input.images, + }; + } else if (promptPayload) { + turn = { + type: 'prompt', + id: input.messageId ?? undefined, + prompt: promptPayload.prompt, + images: input.images, + }; + agent = { + mode: promptPayload.mode, + model: promptPayload.model, + variant: promptPayload.variant, + }; + } else { + throw new Error('sendMessageV2 payload is missing a prompt or command turn'); + } + + const queuedMessage = { + cloudAgentSessionId: input.cloudAgentSessionId, + turn, + agent, + finalization: { + autoCommit: input.autoCommit, + condenseOnComplete: input.condenseOnComplete, + } satisfies TurnFinalization, + ...('completion' in input ? { completion: input.completion } : {}), + }; + const admissionContext = { env: ctx.env, userId: ctx.userId, botId: ctx.botId }; + if (turn.type === 'prompt') { + const replay = await replayMessageIfAlreadyAdmitted(queuedMessage, admissionContext); + if (replay) return withLegacyExecutionId(replay); + + await preflightExistingPromptModel({ + env: ctx.env, + userId: ctx.userId, + cloudAgentSessionId: input.cloudAgentSessionId, + requestedModel: agent?.model, + procedure, + }); + } + + const ack = await queueMessage(queuedMessage, admissionContext); + return withLegacyExecutionId(ack); +} + export function createSessionExecutionV2Handlers() { return { initiateFromKilocodeSessionV2: protectedProcedure @@ -82,70 +166,20 @@ export function createSessionExecutionV2Handlers() { logger.setTags({ userId: ctx.userId, sessionId }); logger.info('Sending V2 message to existing session'); - const commandPayload = - 'payload' in input && input.payload.type === 'command' ? input.payload : undefined; - const promptPayload = - 'prompt' in input - ? { - prompt: input.prompt, - mode: input.mode, - model: input.model, - variant: input.variant, - } - : 'payload' in input && input.payload.type === 'prompt' - ? input.payload - : undefined; - let turn: ExecutionTurnSubmission; - let agent: AgentSelectionOverride | undefined; - if (commandPayload) { - turn = { - type: 'command', - id: input.messageId ?? undefined, - command: commandPayload.command, - arguments: commandPayload.arguments, - images: input.images, - }; - } else if (promptPayload) { - turn = { - type: 'prompt', - id: input.messageId ?? undefined, - prompt: promptPayload.prompt, - images: input.images, - }; - agent = { - mode: promptPayload.mode, - model: promptPayload.model, - variant: promptPayload.variant, - }; - } else { - throw new Error('sendMessageV2 payload is missing a prompt or command turn'); - } - - const queuedMessage = { - cloudAgentSessionId: input.cloudAgentSessionId, - turn, - agent, - finalization: { - autoCommit: input.autoCommit, - condenseOnComplete: input.condenseOnComplete, - } satisfies TurnFinalization, - }; - const admissionContext = { env: ctx.env, userId: ctx.userId, botId: ctx.botId }; - if (turn.type === 'prompt') { - const replay = await replayMessageIfAlreadyAdmitted(queuedMessage, admissionContext); - if (replay) return withLegacyExecutionId(replay); + return queueSendMessageV2(input, ctx, 'sendMessageV2'); + }); + }), - await preflightExistingPromptModel({ - env: ctx.env, - userId: ctx.userId, - cloudAgentSessionId: input.cloudAgentSessionId, - requestedModel: agent?.model, - procedure: 'sendMessageV2', - }); - } + sendMessageV2Internal: internalApiProtectedProcedure + .input(TrustedSendMessageV2Input) + .output(LegacyExecutionResponse) + .mutation(async ({ input, ctx }) => { + return withLogTags({ source: 'sendMessageV2Internal' }, async () => { + const sessionId = input.cloudAgentSessionId as SessionId; + logger.setTags({ userId: ctx.userId, sessionId }); + logger.info('Sending trusted V2 message to existing session'); - const ack = await queueMessage(queuedMessage, admissionContext); - return withLegacyExecutionId(ack); + return queueSendMessageV2(input, ctx, 'sendMessageV2Internal'); }); }), }; diff --git a/services/cloud-agent-next/src/router/handlers/session-management.ts b/services/cloud-agent-next/src/router/handlers/session-management.ts index 2bb86ad23e..750190ead1 100644 --- a/services/cloud-agent-next/src/router/handlers/session-management.ts +++ b/services/cloud-agent-next/src/router/handlers/session-management.ts @@ -25,6 +25,7 @@ import { import { readProfileBundle } from '../../session-profile.js'; import type { CloudAgentSession } from '../../persistence/CloudAgentSession.js'; import type { CloudAgentSessionState } from '../../persistence/types.js'; +import { resolveSessionExecutionPolicy } from '../../persistence/execution-policy.js'; function publicRepositoryFields(metadata: CloudAgentSessionState): { githubRepo?: string; @@ -279,6 +280,7 @@ export function createSessionManagementHandlers() { createdOnPlatform: metadata.identity.createdOnPlatform, appendSystemPrompt: metadata.agent?.appendSystemPrompt, profile: readProfileBundle(metadata), + executionPolicy: resolveSessionExecutionPolicy(metadata), }); // Kill all kilocode processes in this session @@ -641,6 +643,7 @@ export function createSessionManagementHandlers() { profile: sessionService.metadata ? readProfileBundle(sessionService.metadata) : undefined, + executionPolicy: resolveSessionExecutionPolicy(sessionService.metadata), }); // Discover all log files from the sandbox diff --git a/services/cloud-agent-next/src/router/handlers/session-prepare.ts b/services/cloud-agent-next/src/router/handlers/session-prepare.ts index e7ba83b56b..e18b91786e 100644 --- a/services/cloud-agent-next/src/router/handlers/session-prepare.ts +++ b/services/cloud-agent-next/src/router/handlers/session-prepare.ts @@ -258,6 +258,7 @@ export function prepareInputToSessionCreateRequest(input: PrepareInput): Session condenseOnComplete: input.condenseOnComplete, gateThreshold: input.gateThreshold, }, + managedSession: input.managedSession, options: { callbackTarget: input.callbackTarget, kilocodeOrganizationId: input.kilocodeOrganizationId, diff --git a/services/cloud-agent-next/src/router/handlers/session-questions.ts b/services/cloud-agent-next/src/router/handlers/session-questions.ts index 5ae767bcd9..9e676eba65 100644 --- a/services/cloud-agent-next/src/router/handlers/session-questions.ts +++ b/services/cloud-agent-next/src/router/handlers/session-questions.ts @@ -9,12 +9,20 @@ import { protectedProcedure } from '../auth.js'; import { sessionIdSchema } from '../schemas.js'; import { findWrapperForSession } from '../../kilo/wrapper-manager.js'; import { WrapperClient } from '../../kilo/wrapper-client.js'; +import { resolveSessionExecutionPolicy } from '../../persistence/execution-policy.js'; +import { + shouldAutoRejectPermissionInteraction, + shouldDenyQuestionInteraction, +} from '../../shared/managed-session-policy.js'; async function resolveWrapperClient(opts: { sessionId: SessionId; userId: string; env: Env; authToken: string; + interaction?: + | { type: 'question' } + | { type: 'permission'; response: 'once' | 'always' | 'reject' }; }): Promise { const { sessionId, userId, env, authToken } = opts; @@ -22,6 +30,28 @@ async function resolveWrapperClient(opts: { if (!metadata) { throw new TRPCError({ code: 'NOT_FOUND', message: 'Session not found' }); } + const executionPolicy = resolveSessionExecutionPolicy(metadata); + const platform = metadata.identity.createdOnPlatform; + + if ( + opts.interaction?.type === 'question' && + shouldDenyQuestionInteraction({ executionPolicy, platform }) + ) { + throw new TRPCError({ + code: 'FORBIDDEN', + message: 'Questions are disabled for this managed session', + }); + } + if ( + opts.interaction?.type === 'permission' && + opts.interaction.response !== 'reject' && + shouldAutoRejectPermissionInteraction({ executionPolicy, platform }) + ) { + throw new TRPCError({ + code: 'FORBIDDEN', + message: 'Permissions are auto-rejected for this managed session', + }); + } const sandboxId: SandboxId = metadata.workspace?.sandboxId ?? @@ -55,6 +85,9 @@ async function resolveWrapperClient(opts: { env, originalToken: authToken, originalOrgId: metadata.identity.orgId, + createdOnPlatform: metadata.identity.createdOnPlatform, + appendSystemPrompt: metadata.agent?.appendSystemPrompt, + executionPolicy, }); return new WrapperClient({ session, port: wrapperInfo.port }); @@ -82,6 +115,7 @@ export function createSessionQuestionHandlers() { userId, env, authToken: ctx.authToken, + interaction: { type: 'question' }, }); const result = await wrapperClient.answerQuestion(input.questionId, input.answers); logger @@ -119,6 +153,7 @@ export function createSessionQuestionHandlers() { userId, env, authToken: ctx.authToken, + interaction: { type: 'question' }, }); const result = await wrapperClient.rejectQuestion(input.questionId); logger @@ -158,6 +193,7 @@ export function createSessionQuestionHandlers() { userId, env, authToken: ctx.authToken, + interaction: { type: 'permission', response: input.response }, }); const result = await wrapperClient.answerPermission(input.permissionId, input.response); logger diff --git a/services/cloud-agent-next/src/router/schemas.ts b/services/cloud-agent-next/src/router/schemas.ts index 7fa813aae3..c3efa1ac24 100644 --- a/services/cloud-agent-next/src/router/schemas.ts +++ b/services/cloud-agent-next/src/router/schemas.ts @@ -15,6 +15,7 @@ import { RuntimeAgentsSchema, RuntimeKiloCommandsSchema, } from '../persistence/schemas.js'; +import { ManagedSessionExecutionPolicySchema } from '../persistence/execution-policy.js'; import { AgentModeSchema, BUILTIN_AGENT_MODES, Limits } from '../schema.js'; import { MESSAGE_ID_FORMAT_DESCRIPTION, MESSAGE_ID_PATTERN } from '../session/message-id.js'; @@ -218,6 +219,17 @@ const SendMessageV2Options = z.object({ messageId: MessageIdSchema.nullish().describe('Optional message ID for correlating the request'), }); +const MessageCompletionInputSchema = z + .object({ + callbackTarget: CallbackTargetSchema.optional(), + gateThreshold: z.enum(['off', 'all', 'warning', 'critical']).optional(), + }) + .strict(); + +const TrustedSendMessageV2Options = SendMessageV2Options.extend({ + completion: MessageCompletionInputSchema.optional(), +}); + const SendMessageV2FlatInput = SendMessageV2Options.extend(PromptPayload.shape); const SendMessageV2PromptPayloadInput = SendMessageV2Options.extend({ payload: PromptSendPayload, @@ -225,6 +237,13 @@ const SendMessageV2PromptPayloadInput = SendMessageV2Options.extend({ const SendMessageV2CommandPayloadInput = SendMessageV2Options.extend({ payload: CommandSendPayload, }); +const TrustedSendMessageV2FlatInput = TrustedSendMessageV2Options.extend(PromptPayload.shape); +const TrustedSendMessageV2PromptPayloadInput = TrustedSendMessageV2Options.extend({ + payload: PromptSendPayload, +}); +const TrustedSendMessageV2CommandPayloadInput = TrustedSendMessageV2Options.extend({ + payload: CommandSendPayload, +}); export const SendMessageV2Input = z .union([ @@ -251,6 +270,31 @@ export const SendMessageV2Input = z path: ['mode'], }); +export const TrustedSendMessageV2Input = z + .union([ + TrustedSendMessageV2FlatInput, + TrustedSendMessageV2PromptPayloadInput, + TrustedSendMessageV2CommandPayloadInput, + ]) + .transform(input => { + if ('payload' in input && input.payload.type === 'prompt') { + const { payload, ...options } = input; + return { + ...options, + prompt: payload.prompt, + mode: payload.mode, + model: payload.model, + variant: payload.variant, + }; + } + + return input; + }) + .refine(data => !('mode' in data) || data.mode !== 'custom', { + message: 'custom mode requires appendSystemPrompt (use prepareSession/updateSession)', + path: ['mode'], + }); + export type SendMessageV2InputPayload = z.infer; export const PtyIdSchema = z @@ -444,6 +488,13 @@ export const PrepareSessionInput = z .describe( 'PR gate threshold — when not "off", the agent should evaluate findings and report gateResult in its callback' ), + managedSession: z + .object({ + executionPolicy: ManagedSessionExecutionPolicySchema.optional(), + }) + .strict() + .optional() + .describe('Trusted managed-session controls accepted only on internal prepareSession calls'), initialMessageId: MessageIdSchema.optional().describe( 'Initial message ID for correlation with external systems' ), diff --git a/services/cloud-agent-next/src/session-service.test.ts b/services/cloud-agent-next/src/session-service.test.ts index 53f21403b7..951feadc78 100644 --- a/services/cloud-agent-next/src/session-service.test.ts +++ b/services/cloud-agent-next/src/session-service.test.ts @@ -68,6 +68,8 @@ import type { CloudAgentSessionState, PersistenceEnv } from './persistence/types import { parseSessionMetadata } from './persistence/session-metadata.js'; import type { ExecutionSession, SandboxInstance, SessionId } from './types.js'; import type { FencedWrapperDispatchRequest } from './execution/types.js'; +import { buildCodeReviewManagedSessionPolicy } from '@kilocode/worker-utils/managed-session-policy'; +import { ManagedSessionExecutionPolicySchema } from './persistence/execution-policy.js'; type MockExecutionSession = ExecutionSession & { exec: ReturnType; @@ -947,6 +949,41 @@ describe('SessionService.buildWrapperSessionReadyAndPromptRequests', () => { expect(result.promptRequest.session).toEqual(result.readyRequest.session); }); + it('materializes explicit managed session policy into runtime config and wrapper binding', async () => { + const executionPolicy = ManagedSessionExecutionPolicySchema.parse( + buildCodeReviewManagedSessionPolicy() + ); + const metadata = { + ...createMetadata({ createdOnPlatform: 'cloud-agent' }), + executionPolicy, + } satisfies CloudAgentSessionState; + + const result = await buildPromptWrapperRequests(metadata); + + expect(result.readyRequest.materialized.env.CI).toBe('true'); + expect(result.readyRequest.materialized.env.GIT_TERMINAL_PROMPT).toBe('0'); + expect(result.readyRequest.session.executionPolicy).toEqual(executionPolicy); + expect(result.type).toBe('prompt'); + if (result.type !== 'prompt') throw new Error('Expected prompt delivery request'); + expect(result.promptRequest.session.executionPolicy).toEqual(executionPolicy); + + const config = JSON.parse(result.readyRequest.materialized.env.KILO_CONFIG_CONTENT) as { + permission: Record; + }; + expect(config.permission).toMatchObject({ + edit: 'deny', + webfetch: 'deny', + websearch: 'deny', + codesearch: 'deny', + question: 'deny', + }); + expect(config.permission.bash).toMatchObject({ + 'git push': 'deny', + 'git push *': 'deny', + 'gh pr diff': 'allow', + }); + }); + it('materializes OAuth bearer mode with a self-managed GitLab host', async () => { const result = await buildPromptWrapperRequests( createMetadata({ diff --git a/services/cloud-agent-next/src/session-service.ts b/services/cloud-agent-next/src/session-service.ts index 870fcc02fa..9697c3ac58 100644 --- a/services/cloud-agent-next/src/session-service.ts +++ b/services/cloud-agent-next/src/session-service.ts @@ -57,6 +57,10 @@ import { } from './kilo/sandbox-runtime.js'; import { shellQuote, validShellEnvEntries } from './kilo/utils.js'; import { buildSignedImagePromptAttachments } from './execution/image-prompt-parts.js'; +import { + resolveSessionExecutionPolicy, + type ManagedSessionExecutionPolicy, +} from './persistence/execution-policy.js'; import { type WrapperBootstrapRepoSource, type WrapperCommandRequest, @@ -909,6 +913,7 @@ export class SessionService { gitlabTokenManaged: context.gitlabTokenManaged, platform: context.platform, profile: effectiveProfile, + executionPolicy: opts.executionPolicy, }); } @@ -930,6 +935,7 @@ export class SessionService { gitlabTokenManaged, platform, profile, + executionPolicy: storedExecutionPolicy, } = opts; const userEnvVars = profile?.envVars; const encryptedSecrets = profile?.encryptedSecrets; @@ -984,10 +990,15 @@ export class SessionService { if (env.KILO_OPENROUTER_BASE) { providerOptions.baseURL = backendUrlForSandbox(env.KILO_OPENROUTER_BASE); } - const isInteractive = createdOnPlatform == 'cloud-agent-web'; - const commandGuardPolicy = getCommandGuardPolicy(createdOnPlatform); + const managedPolicy = resolveSessionExecutionPolicy({ + executionPolicy: storedExecutionPolicy, + identity: { createdOnPlatform }, + }); + const isInteractive = + createdOnPlatform == 'cloud-agent-web' && managedPolicy?.runtime?.nonInteractive !== true; + const denyQuestions = !isInteractive || managedPolicy?.interaction?.question === 'deny'; - if (commandGuardPolicy) { + if (managedPolicy?.runtime?.nonInteractive === true) { Object.assign(envVars, { CI: 'true', GIT_TERMINAL_PROMPT: '0', @@ -1005,7 +1016,7 @@ export class SessionService { [`${workspacePath}/**`]: 'allow', [`${sessionHome}/.kilocode/skills/**`]: 'allow', }, - ...(!isInteractive && { question: 'deny' }), + ...(denyQuestions && { question: 'deny' }), read: 'allow', edit: 'allow', glob: 'allow', @@ -1023,33 +1034,15 @@ export class SessionService { suggest: 'deny', }; - if (commandGuardPolicy) { - const bashPermissions = buildCommandGuardBashPermissions(commandGuardPolicy); - - // Parity with old autoApproval config: - // read: allow (was read.enabled: true) - // edit: deny (was write.enabled: false) - // webfetch/websearch/codesearch: deny (was browser.enabled: false) - // MCP: allowed by default (was mcp.enabled: true) - // question: handled above (line 564) for non-interactive sessions - Object.assign(permission, { - read: 'allow', - edit: 'deny', - bash: bashPermissions, - webfetch: 'deny', - websearch: 'deny', - codesearch: 'deny', - todowrite: 'allow', - todoread: 'allow', - }); - + if (managedPolicy?.permissionOverrides) { + Object.assign(permission, managedPolicy.permissionOverrides); logger .withFields({ createdOnPlatform, - commandPolicy: commandGuardPolicy.policyName, - deniedCommandPatterns: commandGuardPolicy.denied.length, + executionPolicy: managedPolicy.name, + permissionOverrideKeys: Object.keys(managedPolicy.permissionOverrides).length, }) - .info('Enabled read-only command guard policy'); + .info('Applied managed session execution policy'); } const configContent: Record = { @@ -1182,6 +1175,7 @@ export class SessionService { createdOnPlatform, appendSystemPrompt, profile, + executionPolicy, } = opts; const { sessionId, sessionHome, workspacePath, envVars: contextEnvVars } = context; @@ -1199,6 +1193,7 @@ export class SessionService { createdOnPlatform, appendSystemPrompt, profile: effectiveProfile, + executionPolicy, }); const session = await sandbox.createSession({ @@ -1342,6 +1337,7 @@ export class SessionService { platform, }); + const executionPolicy = resolveSessionExecutionPolicy(metadata); const materializedEnv = this.getSaferEnvVars({ sessionHome, sessionId, @@ -1359,6 +1355,7 @@ export class SessionService { gitlabTokenManaged: resolvedTokens.gitlabTokenManaged, platform, profile, + executionPolicy, }); const ready = { @@ -1382,6 +1379,7 @@ export class SessionService { sessionId, wrapper, upstreamBranch: metadata.repository?.upstreamBranch, + executionPolicy, }); const attachments = @@ -1583,6 +1581,7 @@ export class SessionService { createdOnPlatform: metadata.identity.createdOnPlatform, appendSystemPrompt: metadata.agent?.appendSystemPrompt, profile: readProfileBundle(metadata), + executionPolicy: resolveSessionExecutionPolicy(metadata), }); // Warm fast path: probe for an existing .git before touching disk. The @@ -1784,6 +1783,7 @@ export class SessionService { createdOnPlatform: metadata.identity.createdOnPlatform, appendSystemPrompt: metadata.agent?.appendSystemPrompt, profile: readProfileBundle(metadata), + executionPolicy: resolveSessionExecutionPolicy(metadata), }); } @@ -2390,6 +2390,7 @@ export type GetOrCreateSessionOptions = { createdOnPlatform?: string; appendSystemPrompt?: string; profile?: SessionProfileBundle; + executionPolicy?: ManagedSessionExecutionPolicy; }; export type BuildRuntimeEnvOptions = Omit; @@ -2420,6 +2421,7 @@ type GetSaferEnvVarsOptions = { gitlabTokenManaged?: boolean; platform?: 'github' | 'gitlab'; profile?: SessionProfileBundle; + executionPolicy?: ManagedSessionExecutionPolicy; }; export type WorkspaceReadyMetadata = { @@ -2463,8 +2465,10 @@ function buildWrapperSessionBinding(options: { sessionId: string; wrapper: FencedWrapperDispatchRequest['wrapper']; upstreamBranch?: string; + executionPolicy?: ManagedSessionExecutionPolicy; }): WrapperSessionReadyRequest['session'] { - const { workerUrl, kilocodeToken, userId, sessionId, wrapper, upstreamBranch } = options; + const { workerUrl, kilocodeToken, userId, sessionId, wrapper, upstreamBranch, executionPolicy } = + options; if (!workerUrl) { throw ExecutionError.invalidRequest('WORKER_URL is required for wrapper bootstrap'); } @@ -2480,5 +2484,6 @@ function buildWrapperSessionBinding(options: { wrapperGeneration: wrapper.fence.wrapperGeneration, wrapperConnectionId: wrapper.fence.wrapperConnectionId, ...(upstreamBranch ? { upstreamBranch } : {}), + ...(executionPolicy ? { executionPolicy } : {}), }; } diff --git a/services/cloud-agent-next/src/session/message-settlement-outbox.ts b/services/cloud-agent-next/src/session/message-settlement-outbox.ts index 70f91bb0e6..169ae3512c 100644 --- a/services/cloud-agent-next/src/session/message-settlement-outbox.ts +++ b/services/cloud-agent-next/src/session/message-settlement-outbox.ts @@ -208,7 +208,14 @@ export function createMessageSettlementOutbox( if (representativeMessageId && gateResult !== undefined) { const representative = await getSessionMessageState(storage, representativeMessageId); if (representative?.callbackRequired && !representative.callbackEnqueuedAt) { - await putSessionMessageState(storage, { ...representative, gateResult }); + await putSessionMessageState(storage, { + ...representative, + gateResult, + completionData: { + ...representative.completionData, + gateResult, + }, + }); } } @@ -267,12 +274,19 @@ export function createMessageSettlementOutbox( accepted: true, completionSource: state.completionSource, }; + const completionData = { + ...state.completionData, + ...(extra?.gateResult !== undefined ? { gateResult: extra.gateResult } : {}), + }; if (state.assistantMessageId) { payload.assistantMessageId = state.assistantMessageId; } if (extra?.gateResult !== undefined) { payload.gateResult = extra.gateResult; } + if (Object.keys(completionData).length > 0) { + payload.completionData = completionData; + } ensureTerminalMessageEvent({ entityId: `terminal-message/${state.messageId}`, sessionId, @@ -365,6 +379,23 @@ export function createMessageSettlementOutbox( } } + const completionData = + status === 'completed' + ? { + ...state.completionData, + ...(lastAssistantMessageText !== undefined ? { lastAssistantMessageText } : {}), + ...(metadata?.repository?.upstreamBranch !== undefined + ? { lastSeenBranch: metadata.repository.upstreamBranch } + : {}), + ...(state.gateResult !== undefined ? { gateResult: state.gateResult } : {}), + } + : state.completionData; + + if (status === 'completed' && completionData !== state.completionData) { + state.completionData = completionData; + await putSessionMessageState(storage, state); + } + const payload: CallbackJob['payload'] = { sessionId, cloudAgentSessionId: sessionId, @@ -376,6 +407,7 @@ export function createMessageSettlementOutbox( kiloSessionId: metadata?.auth.kiloSessionId, gateResult: state.gateResult, lastAssistantMessageText, + completionData, idempotencyKey: state.messageId, }; @@ -427,7 +459,9 @@ export function createMessageSettlementOutbox( return false; } - const gateThreshold = (await getMetadata())?.finalization?.gateThreshold; + const gateThreshold = + representative.completionPolicy?.gateThreshold ?? + (await getMetadata())?.finalization?.gateThreshold; return gateThreshold !== undefined && gateThreshold !== 'off'; } diff --git a/services/cloud-agent-next/src/session/pending-messages.ts b/services/cloud-agent-next/src/session/pending-messages.ts index b178258db3..4f62d8d40e 100644 --- a/services/cloud-agent-next/src/session/pending-messages.ts +++ b/services/cloud-agent-next/src/session/pending-messages.ts @@ -57,6 +57,11 @@ const PendingSessionMessageCallbackSnapshotSchema = z target: CallbackTargetSchema.optional(), }) .strict(); +const PendingSessionMessageCompletionPolicySchema = z + .object({ + gateThreshold: z.enum(['off', 'all', 'warning', 'critical']).optional(), + }) + .strict(); const PendingDeliveryDispositionSchema = z.enum(['terminalization-pending']); const PendingDeliverySchema = z.object({ queuedAt: z.number(), @@ -70,6 +75,7 @@ export const PendingSessionMessageV2Schema = z.object({ intent: SessionMessageIntentSchema, delivery: PendingDeliverySchema, callbackSnapshot: PendingSessionMessageCallbackSnapshotSchema.optional(), + completionPolicy: PendingSessionMessageCompletionPolicySchema.optional(), }); export type PendingSessionMessageV2 = z.infer; @@ -124,6 +130,7 @@ export type PendingSessionMessage = { createdAt: number; intent?: SessionMessageIntent; callbackSnapshot?: z.infer; + completionPolicy?: z.infer; flushAttempts?: number; nextFlushAttemptAt?: number; lastFlushError?: string; @@ -249,6 +256,7 @@ function decodePendingMessage( createdAt: message.delivery.queuedAt, intent: message.intent, callbackSnapshot: message.callbackSnapshot, + completionPolicy: message.completionPolicy, flushAttempts: message.delivery.flushAttempts, nextFlushAttemptAt: message.delivery.nextFlushAttemptAt, lastFlushError: message.delivery.lastFlushError, @@ -297,7 +305,8 @@ export function createPendingSessionMessage(params: { export function createPendingSessionMessageFromIntent( intent: SessionMessageIntent, createdAt = Date.now(), - callbackSnapshot?: z.infer + callbackSnapshot?: z.infer, + completionPolicy?: z.infer ): PendingSessionMessage { return { version: 2, @@ -306,6 +315,7 @@ export function createPendingSessionMessageFromIntent( createdAt, intent, callbackSnapshot, + completionPolicy, }; } export function resolvePendingSessionMessageExecutionOptions( @@ -353,6 +363,7 @@ function serializePendingSessionMessage( disposition: normalized.deliveryDisposition, }, callbackSnapshot: normalized.callbackSnapshot, + completionPolicy: normalized.completionPolicy, }) : LegacyPendingSessionMessageSchema.parse({ ...normalized.legacy, diff --git a/services/cloud-agent-next/src/session/queue-message.ts b/services/cloud-agent-next/src/session/queue-message.ts index ba13edac02..d65a65f7f6 100644 --- a/services/cloud-agent-next/src/session/queue-message.ts +++ b/services/cloud-agent-next/src/session/queue-message.ts @@ -130,6 +130,7 @@ export async function queueMessage( }, agent: input.agent, finalization: input.finalization, + completion: input.completion, }; const result = await withDORetry< diff --git a/services/cloud-agent-next/src/session/session-message-queue.test.ts b/services/cloud-agent-next/src/session/session-message-queue.test.ts index cdb5f598ab..6db39109a0 100644 --- a/services/cloud-agent-next/src/session/session-message-queue.test.ts +++ b/services/cloud-agent-next/src/session/session-message-queue.test.ts @@ -27,6 +27,8 @@ import { putSessionMessageState, type TerminalizeParams, } from './session-message-state.js'; +import { buildCodeReviewManagedSessionPolicy } from '@kilocode/worker-utils/managed-session-policy'; +import { ManagedSessionExecutionPolicySchema } from '../persistence/execution-policy.js'; type QueueEvent = { sessionId: string; @@ -524,6 +526,30 @@ describe('SessionMessageQueue', () => { expect(replay).toMatchObject({ success: false, code: 'BAD_REQUEST' }); }); + it('rejects submitted agent overrides for managed sessions', async () => { + const harness = createQueueHarness({ + metadata: createMetadata({ + agent: { mode: 'code-review', model: 'default-model', variant: 'alpha' }, + executionPolicy: ManagedSessionExecutionPolicySchema.parse( + buildCodeReviewManagedSessionPolicy() + ), + }), + }); + + const result = await harness.queue.admitSubmittedMessage({ + userId: 'user_test' as UserId, + turn: { type: 'prompt', id: FIRST_MESSAGE_ID, prompt: 'try another mode' }, + agent: { mode: 'code', model: 'default-model', variant: 'alpha' }, + }); + + expect(result).toMatchObject({ + success: false, + code: 'BAD_REQUEST', + error: 'Managed sessions cannot override agent mode', + }); + await expect(listPendingSessionMessages(harness.storage)).resolves.toHaveLength(0); + }); + it('accepts replay matching predecessor partial immutable constraints without a stored turn', async () => { const harness = createQueueHarness(); await harness.storage.put(`session_message:${FIRST_MESSAGE_ID}`, { diff --git a/services/cloud-agent-next/src/session/session-message-queue.ts b/services/cloud-agent-next/src/session/session-message-queue.ts index 986e7aa80f..6817b97ea0 100644 --- a/services/cloud-agent-next/src/session/session-message-queue.ts +++ b/services/cloud-agent-next/src/session/session-message-queue.ts @@ -2,6 +2,7 @@ import { TRPCError } from '@trpc/server'; import type { ExecutionDeliveryContext, AdmitAcceptedSessionMessageRequest, + MessageCompletionPolicy, MessageDeliveryRequest, MessageDeliveryResult, SessionMessageAdmissionResult, @@ -12,6 +13,7 @@ import { renderExecutionTurnContent } from '../execution/types.js'; import { isExecutionError } from '../execution/errors.js'; import { logger } from '../logger.js'; import { dispatchedKilocodeModelId } from '../persistence/model-utils.js'; +import { resolveSessionExecutionPolicy } from '../persistence/execution-policy.js'; import type { SessionMetadata } from '../persistence/session-metadata.js'; import { isSandboxWorkspaceProbeTimeoutError } from '../sandbox-recovery.js'; import { @@ -41,6 +43,7 @@ import { listReconnectVisibleTerminalQueuedMessages, putSessionMessageState, type SessionMessageStorage, + type SessionMessageCompletionPolicy, type TerminalizeParams, } from './session-message-state.js'; import type { QueuedMessageSnapshot } from '../websocket/stream.js'; @@ -209,6 +212,41 @@ function buildMessageDeliveryRequest( } satisfies MessageDeliveryRequest; } +function completionPolicyFromRequest( + completion: MessageCompletionPolicy | undefined +): SessionMessageCompletionPolicy | undefined { + return completion?.gateThreshold !== undefined + ? { gateThreshold: completion.gateThreshold } + : undefined; +} + +function validateManagedSessionAgentOverride( + metadata: SessionMetadata, + requestedAgent: SubmittedSessionMessageRequest['agent'] +): string | null { + const executionPolicy = resolveSessionExecutionPolicy(metadata); + if (!executionPolicy || !requestedAgent) return null; + + const expectedMode = metadata.agent?.mode ?? 'code'; + if (requestedAgent.mode !== undefined && requestedAgent.mode !== expectedMode) { + return 'Managed sessions cannot override agent mode'; + } + + if (requestedAgent.model !== undefined) { + const expectedModel = dispatchedKilocodeModelId(metadata.agent?.model); + const requestedModel = dispatchedKilocodeModelId(requestedAgent.model); + if (requestedModel !== expectedModel) { + return 'Managed sessions cannot override agent model'; + } + } + + if (requestedAgent.variant !== undefined && requestedAgent.variant !== metadata.agent?.variant) { + return 'Managed sessions cannot override agent variant'; + } + + return null; +} + class MessageDeliveryRequestValidationError extends Error { readonly code = 'BAD_REQUEST' as const; @@ -319,7 +357,12 @@ export async function flushNextPendingSessionMessage(params: { : undefined); await putSessionMessageState( params.storage, - createQueuedSessionMessageState(intent, callbackSnapshot, message.createdAt) + createQueuedSessionMessageState( + intent, + callbackSnapshot, + message.createdAt, + message.completionPolicy + ) ); } await params.repairQueuedMessageEffects?.(intent); @@ -599,11 +642,19 @@ export function createSessionMessageQueue( : undefined); await putSessionMessageState( storage, - createQueuedSessionMessageState(intent, callbackSnapshot, message.createdAt) + createQueuedSessionMessageState( + intent, + callbackSnapshot, + message.createdAt, + message.completionPolicy + ) ); } - async function admitIntent(intent: SessionMessageIntent): Promise { + async function admitIntent( + intent: SessionMessageIntent, + completion?: MessageCompletionPolicy + ): Promise { const { turn } = intent; const idempotentResult = await getExistingAdmissionAckForMessageId(turn.messageId, intent); if (idempotentResult) return idempotentResult; @@ -612,13 +663,25 @@ export function createSessionMessageQueue( if (capacityError) return capacityError; const metadata = await getMetadata(); - const callbackTarget = metadata?.callback?.target; + const callbackTarget = completion?.callbackTarget ?? metadata?.callback?.target; const callbackSnapshot = callbackTarget ? { required: true, target: callbackTarget } : undefined; + const completionPolicy = completionPolicyFromRequest(completion); - await enqueuePendingSessionMessageIntent(storage, intent, Date.now(), callbackSnapshot); - const messageState = createQueuedSessionMessageState(intent, callbackSnapshot); + await enqueuePendingSessionMessageIntent( + storage, + intent, + Date.now(), + callbackSnapshot, + completionPolicy + ); + const messageState = createQueuedSessionMessageState( + intent, + callbackSnapshot, + Date.now(), + completionPolicy + ); await putSessionMessageState(storage, messageState); await completeQueuedAdmissionEffects(intent); return buildAdmissionAck(turn.messageId); @@ -693,6 +756,13 @@ export function createSessionMessageQueue( const requestedAgent = request.agent; const requestedFinalization = request.finalization; + const managedAgentOverrideError = validateManagedSessionAgentOverride( + metadata, + requestedAgent + ); + if (managedAgentOverrideError) { + return buildAdmissionError('BAD_REQUEST', managedAgentOverrideError); + } const modeInput = requestedAgent?.mode ?? metadata.agent?.mode ?? 'code'; const modeCheck = validateModeAgainstRuntimeAgents(metadata, modeInput); if (modeCheck) { @@ -737,7 +807,7 @@ export function createSessionMessageQueue( if (existingAdmission) return existingAdmission; const capacityError = await checkPendingQueueCapacity(); if (capacityError) return capacityError; - return await admitIntent(intent); + return await admitIntent(intent, request.completion); } catch (error) { if (isExecutionError(error)) { if (error.retryable) { @@ -976,9 +1046,15 @@ export async function enqueuePendingSessionMessageIntent( storage: SessionQueueStorage, intent: SessionMessageIntent, createdAt = Date.now(), - callbackSnapshot?: PendingSessionMessage['callbackSnapshot'] + callbackSnapshot?: PendingSessionMessage['callbackSnapshot'], + completionPolicy?: PendingSessionMessage['completionPolicy'] ): Promise { - const message = createPendingSessionMessageFromIntent(intent, createdAt, callbackSnapshot); + const message = createPendingSessionMessageFromIntent( + intent, + createdAt, + callbackSnapshot, + completionPolicy + ); await storePendingSessionMessage(storage, message); return message; } diff --git a/services/cloud-agent-next/src/session/session-message-state.ts b/services/cloud-agent-next/src/session/session-message-state.ts index 0eba922339..8e9146ea2d 100644 --- a/services/cloud-agent-next/src/session/session-message-state.ts +++ b/services/cloud-agent-next/src/session/session-message-state.ts @@ -35,6 +35,19 @@ export type TerminalEffectAccounting = { callback: TerminalCallbackEffectAccounting; }; +export type SessionMessageCompletionPolicy = { + gateThreshold?: 'off' | 'all' | 'warning' | 'critical'; +}; + +export type SessionMessageCompletionData = { + assistantMessageId?: string; + lastAssistantMessageText?: string; + gateResult?: 'pass' | 'fail'; + lastSeenBranch?: string; + completedAt?: number; + source?: SessionMessageCompletionSource; +}; + export type SessionMessageState = { messageId: string; status: SessionMessageStatus; @@ -53,6 +66,8 @@ export type SessionMessageState = { failureReason?: string; attempts?: number; gateResult?: 'pass' | 'fail'; + completionPolicy?: SessionMessageCompletionPolicy; + completionData?: SessionMessageCompletionData; callbackRequired?: boolean; callbackTarget?: CallbackTarget; callbackEnqueuedAt?: number; @@ -156,6 +171,31 @@ export const SessionMessageStateSchema = z failureReason: z.string().optional(), attempts: z.number().int().nonnegative().optional(), gateResult: z.enum(['pass', 'fail']).optional(), + completionPolicy: z + .object({ + gateThreshold: z.enum(['off', 'all', 'warning', 'critical']).optional(), + }) + .strict() + .optional(), + completionData: z + .object({ + assistantMessageId: z.string().optional(), + lastAssistantMessageText: z.string().optional(), + gateResult: z.enum(['pass', 'fail']).optional(), + lastSeenBranch: z.string().optional(), + completedAt: z.number().optional(), + source: z + .enum([ + 'assistant_message_event', + 'idle_reconciliation', + 'wrapper_failure', + 'interrupt', + 'delivery_failure', + ]) + .optional(), + }) + .strict() + .optional(), callbackRequired: z.boolean().optional(), callbackTarget: z .object({ @@ -280,7 +320,8 @@ export async function putSessionMessageState( export function createQueuedSessionMessageState( intent: SessionMessageIntent, callbackSnapshot?: { required: boolean; target?: CallbackTarget }, - now = Date.now() + now = Date.now(), + completionPolicy?: SessionMessageCompletionPolicy ): SessionMessageState { return { messageId: intent.turn.messageId, @@ -291,6 +332,7 @@ export function createQueuedSessionMessageState( queuedAt: now, callbackRequired: callbackSnapshot?.required ?? false, callbackTarget: callbackSnapshot?.target, + completionPolicy, }; } @@ -317,6 +359,7 @@ export type MarkMessageCompletedParams = { assistantMessageId?: string; completionSource: SessionMessageCompletionSource; gateResult?: 'pass' | 'fail'; + completionData?: SessionMessageCompletionData; }; export async function markMessageCompleted( @@ -328,13 +371,28 @@ export async function markMessageCompleted( const state = await getSessionMessageState(storage, messageId); if (!state) return null; if (isTerminalStatus(state.status)) return null; + const assistantMessageId = + params.assistantMessageId ?? + params.completionData?.assistantMessageId ?? + state.completionData?.assistantMessageId ?? + state.assistantMessageId; + const gateResult = + params.gateResult ?? params.completionData?.gateResult ?? state.completionData?.gateResult; const updated: SessionMessageState = { ...state, status: 'completed', terminalAt: now, - assistantMessageId: params.assistantMessageId, + assistantMessageId, completionSource: params.completionSource, - gateResult: params.gateResult, + gateResult, + completionData: { + ...state.completionData, + ...params.completionData, + ...(assistantMessageId !== undefined ? { assistantMessageId } : {}), + ...(gateResult !== undefined ? { gateResult } : {}), + completedAt: now, + source: params.completionSource, + }, }; await putSessionMessageState(storage, updated); return updated; @@ -501,6 +559,7 @@ export type TerminalizeParams = assistantMessageId?: string; completionSource: SessionMessageCompletionSource; gateResult?: 'pass' | 'fail'; + completionData?: SessionMessageCompletionData; } | { kind: 'failed'; @@ -537,13 +596,29 @@ export async function terminalizeMessageOnce( }; let updated: SessionMessageState; if (params.kind === 'completed') { + const assistantMessageId = + params.assistantMessageId ?? + params.completionData?.assistantMessageId ?? + state.completionData?.assistantMessageId ?? + state.assistantMessageId; + const gateResult = + params.gateResult ?? params.completionData?.gateResult ?? state.completionData?.gateResult; + const completionData: SessionMessageCompletionData = { + ...state.completionData, + ...params.completionData, + ...(assistantMessageId !== undefined ? { assistantMessageId } : {}), + ...(gateResult !== undefined ? { gateResult } : {}), + completedAt: now, + source: params.completionSource, + }; updated = { ...state, status: 'completed', terminalAt: now, - assistantMessageId: params.assistantMessageId, + assistantMessageId, completionSource: params.completionSource, - gateResult: params.gateResult, + gateResult, + completionData, terminalEffects, }; } else if (params.kind === 'failed') { diff --git a/services/cloud-agent-next/src/session/session-registration.ts b/services/cloud-agent-next/src/session/session-registration.ts index 762facc53c..626e7127a8 100644 --- a/services/cloud-agent-next/src/session/session-registration.ts +++ b/services/cloud-agent-next/src/session/session-registration.ts @@ -178,6 +178,7 @@ function buildSessionRegistrationCommand( repository: input.repository, profile: input.profile?.resolved, finalization: input.finalization, + executionPolicy: input.managedSession?.executionPolicy, callback: input.options?.callbackTarget ? { target: input.options.callbackTarget } : undefined, workspace: { sandboxId: allocation.sandboxId, diff --git a/services/cloud-agent-next/src/session/session-requests.ts b/services/cloud-agent-next/src/session/session-requests.ts index 5d77246ee2..4ca5c58d42 100644 --- a/services/cloud-agent-next/src/session/session-requests.ts +++ b/services/cloud-agent-next/src/session/session-requests.ts @@ -1,4 +1,5 @@ import type { CallbackTarget } from '../callbacks/index.js'; +import type { ManagedSessionExecutionPolicy } from '../persistence/execution-policy.js'; import type { AgentSelection, ExecutionTurnSubmission, @@ -49,6 +50,9 @@ export type SessionCreateRequest = { resolved?: SessionProfileBundle; }; finalization?: SessionFinalization; + managedSession?: { + executionPolicy?: ManagedSessionExecutionPolicy; + }; options?: { callbackTarget?: CallbackTarget; kilocodeOrganizationId?: string; diff --git a/services/cloud-agent-next/src/shared/managed-session-policy.ts b/services/cloud-agent-next/src/shared/managed-session-policy.ts new file mode 100644 index 0000000000..a9719f33d0 --- /dev/null +++ b/services/cloud-agent-next/src/shared/managed-session-policy.ts @@ -0,0 +1,88 @@ +export type ManagedSessionInteractionPolicy = { + question?: 'allow' | 'deny'; + permission?: 'allow' | 'auto-approve' | 'auto-reject'; + terminal?: 'allow' | 'deny'; +}; + +export type ManagedSessionExecutionPolicy = { + name: string; + permissionOverrides?: Record; + tools?: Record; + interaction?: ManagedSessionInteractionPolicy; + runtime?: { + nonInteractive?: boolean; + }; +}; + +export const LEGACY_CODE_REVIEW_INTERACTION_POLICY = { + question: 'deny', + permission: 'auto-reject', + terminal: 'deny', +} satisfies ManagedSessionInteractionPolicy; + +function isRecord(value: unknown): value is Record { + return typeof value === 'object' && value !== null && !Array.isArray(value); +} + +function isInteractionPolicy(value: unknown): value is ManagedSessionInteractionPolicy { + if (!isRecord(value)) return false; + const question = value.question; + const permission = value.permission; + const terminal = value.terminal; + return ( + (question === undefined || question === 'allow' || question === 'deny') && + (permission === undefined || + permission === 'allow' || + permission === 'auto-approve' || + permission === 'auto-reject') && + (terminal === undefined || terminal === 'allow' || terminal === 'deny') + ); +} + +export function isManagedSessionExecutionPolicy( + value: unknown +): value is ManagedSessionExecutionPolicy { + if (!isRecord(value)) return false; + if (typeof value.name !== 'string' || value.name.length === 0) return false; + if (value.permissionOverrides !== undefined && !isRecord(value.permissionOverrides)) return false; + if (value.tools !== undefined) { + if (!isRecord(value.tools)) return false; + if (!Object.values(value.tools).every(item => typeof item === 'boolean')) return false; + } + if (value.interaction !== undefined && !isInteractionPolicy(value.interaction)) return false; + if (value.runtime !== undefined) { + if (!isRecord(value.runtime)) return false; + const nonInteractive = value.runtime.nonInteractive; + if (nonInteractive !== undefined && typeof nonInteractive !== 'boolean') return false; + } + return true; +} + +export function effectiveInteractionPolicy(input: { + executionPolicy?: ManagedSessionExecutionPolicy; + platform?: string; +}): ManagedSessionInteractionPolicy | undefined { + if (input.executionPolicy?.interaction) return input.executionPolicy.interaction; + return input.platform === 'code-review' ? LEGACY_CODE_REVIEW_INTERACTION_POLICY : undefined; +} + +export function shouldDenyQuestionInteraction(input: { + executionPolicy?: ManagedSessionExecutionPolicy; + platform?: string; +}): boolean { + return effectiveInteractionPolicy(input)?.question === 'deny'; +} + +export function shouldAutoRejectPermissionInteraction(input: { + executionPolicy?: ManagedSessionExecutionPolicy; + platform?: string; +}): boolean { + return effectiveInteractionPolicy(input)?.permission === 'auto-reject'; +} + +export function shouldDenyTerminalInteraction(input: { + executionPolicy?: ManagedSessionExecutionPolicy; + platform?: string; +}): boolean { + return effectiveInteractionPolicy(input)?.terminal === 'deny'; +} diff --git a/services/cloud-agent-next/src/shared/wrapper-bootstrap.ts b/services/cloud-agent-next/src/shared/wrapper-bootstrap.ts index 884687face..395aa5bb44 100644 --- a/services/cloud-agent-next/src/shared/wrapper-bootstrap.ts +++ b/services/cloud-agent-next/src/shared/wrapper-bootstrap.ts @@ -1,3 +1,5 @@ +import type { ManagedSessionExecutionPolicy } from './managed-session-policy.js'; + export type WrapperBootstrapRepoSource = | { kind: 'github'; @@ -64,6 +66,7 @@ export type WrapperSessionBinding = { ingestToken?: string; workerAuthToken: string; upstreamBranch?: string; + executionPolicy?: ManagedSessionExecutionPolicy; wrapperRunId: string; wrapperGeneration: number; wrapperConnectionId: string; diff --git a/services/cloud-agent-next/src/terminal/access.test.ts b/services/cloud-agent-next/src/terminal/access.test.ts index c14515ea1a..f71e7a012c 100644 --- a/services/cloud-agent-next/src/terminal/access.test.ts +++ b/services/cloud-agent-next/src/terminal/access.test.ts @@ -3,6 +3,8 @@ import type { CloudAgentSessionState } from '../persistence/types.js'; import { WRAPPER_VERSION } from '../shared/wrapper-version.js'; import type { Env, SandboxInstance } from '../types.js'; import { resolveTerminalWrapperClient, validateTerminalMetadata } from './access.js'; +import { buildCodeReviewManagedSessionPolicy } from '@kilocode/worker-utils/managed-session-policy'; +import { ManagedSessionExecutionPolicySchema } from '../persistence/execution-policy.js'; vi.mock('@cloudflare/sandbox', () => ({ getSandbox: vi.fn(), @@ -54,7 +56,24 @@ describe('validateTerminalMetadata', () => { expect(result).toEqual({ success: false, - error: 'Terminal is only available for interactive Cloud Agent sessions', + error: 'Terminal is disabled for this managed session', + }); + }); + + it('rejects explicitly managed sessions even from interactive platforms', () => { + const result = validateTerminalMetadata( + { + ...baseMetadata, + executionPolicy: ManagedSessionExecutionPolicySchema.parse( + buildCodeReviewManagedSessionPolicy() + ), + }, + baseMetadata.identity.sessionId + ); + + expect(result).toEqual({ + success: false, + error: 'Terminal is disabled for this managed session', }); }); diff --git a/services/cloud-agent-next/src/terminal/access.ts b/services/cloud-agent-next/src/terminal/access.ts index 120343c796..900c71c5a0 100644 --- a/services/cloud-agent-next/src/terminal/access.ts +++ b/services/cloud-agent-next/src/terminal/access.ts @@ -8,6 +8,8 @@ import { } from '../kilo/wrapper-client.js'; import { findWrapperForSession } from '../kilo/wrapper-manager.js'; import { generateSandboxId, getSandboxNamespace } from '../sandbox-id.js'; +import { resolveSessionExecutionPolicy } from '../persistence/execution-policy.js'; +import { shouldDenyTerminalInteraction } from '../shared/managed-session-policy.js'; import { WRAPPER_VERSION } from '../shared/wrapper-version.js'; import type { Env, SandboxId, SandboxInstance } from '../types.js'; @@ -36,6 +38,19 @@ export function validateTerminalMetadata( }; } + const executionPolicy = resolveSessionExecutionPolicy(metadata); + if ( + shouldDenyTerminalInteraction({ + executionPolicy, + platform: metadata.identity.createdOnPlatform, + }) + ) { + return { + success: false, + error: 'Terminal is disabled for this managed session', + }; + } + if (!isTerminalSessionPlatform(metadata.identity.createdOnPlatform)) { return { success: false, diff --git a/services/cloud-agent-next/test/unit/wrapper/server.test.ts b/services/cloud-agent-next/test/unit/wrapper/server.test.ts index 54b0c1c142..b15b0b50cd 100644 --- a/services/cloud-agent-next/test/unit/wrapper/server.test.ts +++ b/services/cloud-agent-next/test/unit/wrapper/server.test.ts @@ -19,6 +19,7 @@ import { type SessionBinding, } from '../../../wrapper/src/server.js'; import type { WrapperKiloClient } from '../../../wrapper/src/kilo-api.js'; +import { buildCodeReviewManagedSessionPolicy } from '@kilocode/worker-utils/managed-session-policy'; // --------------------------------------------------------------------------- // Test Helpers @@ -103,6 +104,29 @@ describe('createAnswerPermissionHandler', () => { expect(data).toEqual({ status: 'answered', success: true }); expect(deps.kiloClient.answerPermission).toHaveBeenCalledWith('perm_1', 'always'); }); + + it('rejects permission approvals for managed sessions', async () => { + const state = new WrapperState(); + state.bindSession({ + kiloSessionId: 'kilo_sess_1', + ingestUrl: 'wss://ingest.example.com', + ingestToken: 'token', + workerAuthToken: 'auth', + executionPolicy: buildCodeReviewManagedSessionPolicy(), + }); + const deps = createMockDeps(state); + const handler = createAnswerPermissionHandler(deps); + + const response = await handler(jsonRequest({ permissionId: 'perm_1', response: 'always' })); + const data = await readJson(response); + + expect(response.status).toBe(403); + expect(data).toEqual({ + error: 'PERMISSION_DISABLED', + message: 'Permissions are auto-rejected for this managed session', + }); + expect(deps.kiloClient.answerPermission).not.toHaveBeenCalled(); + }); }); // --------------------------------------------------------------------------- @@ -140,6 +164,29 @@ describe('createAnswerQuestionHandler', () => { expect(data).toEqual({ status: 'answered', success: true }); expect(deps.kiloClient.answerQuestion).toHaveBeenCalledWith('q_1', [['yes']]); }); + + it('rejects question answers for managed sessions', async () => { + const state = new WrapperState(); + state.bindSession({ + kiloSessionId: 'kilo_sess_1', + ingestUrl: 'wss://ingest.example.com', + ingestToken: 'token', + workerAuthToken: 'auth', + executionPolicy: buildCodeReviewManagedSessionPolicy(), + }); + const deps = createMockDeps(state); + const handler = createAnswerQuestionHandler(deps); + + const response = await handler(jsonRequest({ questionId: 'q_1', answers: [['yes']] })); + const data = await readJson(response); + + expect(response.status).toBe(403); + expect(data).toEqual({ + error: 'QUESTION_DISABLED', + message: 'Questions are disabled for this managed session', + }); + expect(deps.kiloClient.answerQuestion).not.toHaveBeenCalled(); + }); }); // --------------------------------------------------------------------------- @@ -177,6 +224,29 @@ describe('createRejectQuestionHandler', () => { expect(data).toEqual({ status: 'rejected', success: true }); expect(deps.kiloClient.rejectQuestion).toHaveBeenCalledWith('q_1'); }); + + it('rejects question rejection endpoint for managed sessions', async () => { + const state = new WrapperState(); + state.bindSession({ + kiloSessionId: 'kilo_sess_1', + ingestUrl: 'wss://ingest.example.com', + ingestToken: 'token', + workerAuthToken: 'auth', + executionPolicy: buildCodeReviewManagedSessionPolicy(), + }); + const deps = createMockDeps(state); + const handler = createRejectQuestionHandler(deps); + + const response = await handler(jsonRequest({ questionId: 'q_1' })); + const data = await readJson(response); + + expect(response.status).toBe(403); + expect(data).toEqual({ + error: 'QUESTION_DISABLED', + message: 'Questions are disabled for this managed session', + }); + expect(deps.kiloClient.rejectQuestion).not.toHaveBeenCalled(); + }); }); // --------------------------------------------------------------------------- diff --git a/services/cloud-agent-next/wrapper/src/connection.ts b/services/cloud-agent-next/wrapper/src/connection.ts index 5766ba3b6a..43288b7616 100644 --- a/services/cloud-agent-next/wrapper/src/connection.ts +++ b/services/cloud-agent-next/wrapper/src/connection.ts @@ -11,6 +11,10 @@ import type { WrapperState } from './state.js'; import type { IngestEvent, WrapperCommand } from '../../src/shared/protocol.js'; +import { + shouldAutoRejectPermissionInteraction, + shouldDenyQuestionInteraction, +} from '../../src/shared/managed-session-policy.js'; import { trimPayload } from '../../src/shared/trim-payload.js'; import { logToFile } from './utils.js'; import type { KiloEvent, WrapperKiloClient } from './kilo-api.js'; @@ -19,8 +23,29 @@ function isRecord(value: unknown): value is Record { return typeof value === 'object' && value !== null; } -function isCodeReviewJob(state: WrapperState): boolean { - return state.currentSession?.platform === 'code-review'; +function shouldRejectPermission(state: WrapperState): boolean { + const session = state.currentSession; + return shouldAutoRejectPermissionInteraction({ + executionPolicy: session?.executionPolicy, + platform: session?.platform, + }); +} + +function shouldRejectQuestion(state: WrapperState): boolean { + const session = state.currentSession; + return shouldDenyQuestionInteraction({ + executionPolicy: session?.executionPolicy, + platform: session?.platform, + }); +} + +function shouldSuppressInteractiveStatus( + state: WrapperState, + statusType: string | undefined +): boolean { + if (statusType === 'question') return shouldRejectQuestion(state); + if (statusType === 'permission') return shouldRejectPermission(state); + return false; } function gateResultFromProperties( @@ -35,10 +60,6 @@ function statusTypeFromProperties(properties: Record): string | return isRecord(status) && typeof status.type === 'string' ? status.type : undefined; } -function isInteractiveStatusType(statusType: string | undefined): boolean { - return statusType === 'question' || statusType === 'permission'; -} - function permissionCategoryFromProperties(properties: Record): string { const permission = properties.permission; if (isRecord(permission)) { @@ -537,11 +558,12 @@ export function createConnectionManager( const pendingQuestion = questions.find(q => q.sessionID === kiloSessionId); const pendingPermission = permissions.find(p => p.sessionID === kiloSessionId); const pendingNetworkWaits = networkWaits.filter(wait => wait.sessionID === kiloSessionId); - const codeReviewJob = isCodeReviewJob(state); - const skipStatusForCodeReview = codeReviewJob && isInteractiveStatusType(sessionStatus.type); + const suppressQuestion = shouldRejectQuestion(state); + const suppressPermission = shouldRejectPermission(state); + const skipInteractiveStatus = shouldSuppressInteractiveStatus(state, sessionStatus.type); // Send session status as a regular kilocode event - if (!skipStatusForCodeReview) { + if (!skipInteractiveStatus) { const statusProperties = { sessionID: kiloSessionId, status: sessionStatus }; sendToIngest({ streamEventType: 'kilocode', @@ -557,7 +579,7 @@ export function createConnectionManager( // Replay pending questions/permissions as regular events // (same format as real-time delivery - matches CLI behavior) - if (pendingQuestion && !codeReviewJob) { + if (pendingQuestion && !suppressQuestion) { sendToIngest({ streamEventType: 'kilocode', data: { @@ -568,7 +590,7 @@ export function createConnectionManager( timestamp: new Date().toISOString(), }); } - if (pendingPermission && !codeReviewJob) { + if (pendingPermission && !suppressPermission) { sendToIngest({ streamEventType: 'kilocode', data: { @@ -593,7 +615,7 @@ export function createConnectionManager( } logToFile( - `kilo state sent: status=${sessionStatus.type}${skipStatusForCodeReview ? ' (suppressed)' : ''}, question=${pendingQuestion?.id ?? 'none'}${codeReviewJob && pendingQuestion ? ' (suppressed)' : ''}, permission=${pendingPermission?.id ?? 'none'}${codeReviewJob && pendingPermission ? ' (suppressed)' : ''}, networkWaits=${pendingNetworkWaits.length}` + `kilo state sent: status=${sessionStatus.type}${skipInteractiveStatus ? ' (suppressed)' : ''}, question=${pendingQuestion?.id ?? 'none'}${suppressQuestion && pendingQuestion ? ' (suppressed)' : ''}, permission=${pendingPermission?.id ?? 'none'}${suppressPermission && pendingPermission ? ' (suppressed)' : ''}, networkWaits=${pendingNetworkWaits.length}` ); } catch (err) { logToFile( @@ -962,7 +984,7 @@ export function createConnectionManager( // waiting for a human response that will never come. if (eventType === 'permission.asked') { const permId = typeof properties.id === 'string' ? properties.id : undefined; - if (isCodeReviewJob(state)) { + if (shouldRejectPermission(state)) { rejectCodeReviewPermission(permId, properties, state, config.kiloClient); callbacks.onSseEvent?.(); continue; @@ -980,21 +1002,19 @@ export function createConnectionManager( continue; } - if (isCodeReviewJob(state)) { - if (eventType === 'question.asked') { - const questionId = typeof properties.id === 'string' ? properties.id : undefined; - rejectCodeReviewQuestion(questionId, config.kiloClient); - callbacks.onSseEvent?.(); - continue; - } + if (eventType === 'question.asked' && shouldRejectQuestion(state)) { + const questionId = typeof properties.id === 'string' ? properties.id : undefined; + rejectCodeReviewQuestion(questionId, config.kiloClient); + callbacks.onSseEvent?.(); + continue; + } - if ( - eventType === 'session.status' && - isInteractiveStatusType(statusTypeFromProperties(properties)) - ) { - callbacks.onSseEvent?.(); - continue; - } + if ( + eventType === 'session.status' && + shouldSuppressInteractiveStatus(state, statusTypeFromProperties(properties)) + ) { + callbacks.onSseEvent?.(); + continue; } maybeResumeNetworkWait(eventType, properties); diff --git a/services/cloud-agent-next/wrapper/src/server.ts b/services/cloud-agent-next/wrapper/src/server.ts index 19281af47b..0773580046 100644 --- a/services/cloud-agent-next/wrapper/src/server.ts +++ b/services/cloud-agent-next/wrapper/src/server.ts @@ -25,6 +25,11 @@ import { type WrapperSessionReadyRequest, type WrapperSessionReadyResponse, } from '../../src/shared/wrapper-bootstrap.js'; +import { + shouldAutoRejectPermissionInteraction, + shouldDenyQuestionInteraction, + type ManagedSessionExecutionPolicy, +} from '../../src/shared/managed-session-policy.js'; // --------------------------------------------------------------------------- // Types @@ -67,6 +72,7 @@ export type SessionBinding = { ingestToken?: string; workerAuthToken: string; upstreamBranch?: string; + executionPolicy?: ManagedSessionExecutionPolicy; wrapperRunId: string; wrapperGeneration: number; wrapperConnectionId: string; @@ -239,6 +245,7 @@ export async function bindSessionContext( ingestToken: binding.ingestToken, workerAuthToken: binding.workerAuthToken, platform: config.platform, + executionPolicy: binding.executionPolicy, wrapperRunId: binding.wrapperRunId, wrapperGeneration: binding.wrapperGeneration, wrapperConnectionId: binding.wrapperConnectionId, @@ -269,6 +276,7 @@ export async function bindSessionContext( ingestToken: binding.ingestToken, workerAuthToken: binding.workerAuthToken, platform: config.platform, + executionPolicy: binding.executionPolicy, wrapperRunId: binding.wrapperRunId, wrapperGeneration: binding.wrapperGeneration, wrapperConnectionId: binding.wrapperConnectionId, @@ -479,6 +487,21 @@ export function createAnswerPermissionHandler(deps: ServerDependencies) { return errorResponse('INVALID_REQUEST', 'permissionId and response are required', 400); } + const session = state.currentSession; + if ( + body.response !== 'reject' && + shouldAutoRejectPermissionInteraction({ + executionPolicy: session?.executionPolicy, + platform: session?.platform, + }) + ) { + return errorResponse( + 'PERMISSION_DISABLED', + 'Permissions are auto-rejected for this managed session', + 403 + ); + } + try { const success = body.message === undefined @@ -516,6 +539,20 @@ export function createAnswerQuestionHandler(deps: ServerDependencies) { return errorResponse('INVALID_REQUEST', 'questionId and answers are required', 400); } + const session = state.currentSession; + if ( + shouldDenyQuestionInteraction({ + executionPolicy: session?.executionPolicy, + platform: session?.platform, + }) + ) { + return errorResponse( + 'QUESTION_DISABLED', + 'Questions are disabled for this managed session', + 403 + ); + } + try { const success = await kiloClient.answerQuestion(body.questionId, body.answers); state.updateActivity(); @@ -548,6 +585,20 @@ export function createRejectQuestionHandler(deps: ServerDependencies) { return errorResponse('INVALID_REQUEST', 'questionId is required', 400); } + const session = state.currentSession; + if ( + shouldDenyQuestionInteraction({ + executionPolicy: session?.executionPolicy, + platform: session?.platform, + }) + ) { + return errorResponse( + 'QUESTION_DISABLED', + 'Questions are disabled for this managed session', + 403 + ); + } + try { const success = await kiloClient.rejectQuestion(body.questionId); state.updateActivity(); diff --git a/services/cloud-agent-next/wrapper/src/state.ts b/services/cloud-agent-next/wrapper/src/state.ts index 216f8bd34a..8a9ae53b9e 100644 --- a/services/cloud-agent-next/wrapper/src/state.ts +++ b/services/cloud-agent-next/wrapper/src/state.ts @@ -13,6 +13,7 @@ */ import type { IngestEvent } from '../../src/shared/protocol.js'; +import type { ManagedSessionExecutionPolicy } from '../../src/shared/managed-session-policy.js'; import type { LogUploader } from './log-uploader.js'; export type { LogUploader } from './log-uploader.js'; @@ -26,6 +27,7 @@ export type SessionContext = { ingestToken?: string; workerAuthToken: string; platform?: string; + executionPolicy?: ManagedSessionExecutionPolicy; wrapperRunId?: string; wrapperGeneration?: number; wrapperConnectionId?: string; @@ -43,6 +45,13 @@ export type MessageInfo = { upstreamBranch?: string; }; +function sameExecutionPolicy( + left: ManagedSessionExecutionPolicy | undefined, + right: ManagedSessionExecutionPolicy | undefined +): boolean { + return JSON.stringify(left) === JSON.stringify(right); +} + export type LastError = { code: string; messageId?: string; @@ -306,6 +315,7 @@ export class WrapperState { this.session.ingestToken !== context.ingestToken || this.session.workerAuthToken !== context.workerAuthToken || this.session.platform !== context.platform || + !sameExecutionPolicy(this.session.executionPolicy, context.executionPolicy) || this.session.wrapperRunId !== context.wrapperRunId || this.session.wrapperGeneration !== context.wrapperGeneration || this.session.wrapperConnectionId !== context.wrapperConnectionId; From 997ce5c9b70ca3e58a6bd696ccd880d0bcec139d Mon Sep 17 00:00:00 2001 From: Alex Alecu Date: Wed, 27 May 2026 21:41:14 +0300 Subject: [PATCH 3/4] feat(code-review): use managed runtime agent --- .../src/code-review-orchestrator.ts | 66 +++++---- .../code-review-orchestrator.test.ts | 132 ++++++++++++------ 2 files changed, 133 insertions(+), 65 deletions(-) diff --git a/services/code-review-infra/src/code-review-orchestrator.ts b/services/code-review-infra/src/code-review-orchestrator.ts index 30a547a131..48350cdc83 100644 --- a/services/code-review-infra/src/code-review-orchestrator.ts +++ b/services/code-review-infra/src/code-review-orchestrator.ts @@ -11,7 +11,10 @@ import { createCloudAgentNextFetchClient, CloudAgentNextBillingError, CloudAgentNextError, + CODE_REVIEW_RUNTIME_AGENT_SLUG, deriveCallbackToken, + buildCodeReviewManagedSessionPolicy, + buildCodeReviewRuntimeAgent, type CloudAgentNextFetchClient, type CloudAgentSessionHealthOutput, type CloudAgentTerminalReason, @@ -1139,10 +1142,22 @@ export class CodeReviewOrchestrator extends DurableObject { this.env.CALLBACK_TOKEN_SECRET ); + const executionPolicy = buildCodeReviewManagedSessionPolicy(); + const runtimeAgent = buildCodeReviewRuntimeAgent({ + model: this.state.sessionInput.model, + ...(this.state.sessionInput.variant ? { variant: this.state.sessionInput.variant } : {}), + ...(executionPolicy.permissionOverrides + ? { permission: executionPolicy.permissionOverrides } + : {}), + }); + const prepareInput = { ...this.state.sessionInput, + mode: CODE_REVIEW_RUNTIME_AGENT_SLUG, createdOnPlatform: 'code-review' as const, callbackTarget, + runtimeAgents: [runtimeAgent], + managedSession: { executionPolicy }, }; console.log('[CodeReviewOrchestrator] Calling prepareSession', { @@ -1266,13 +1281,13 @@ export class CodeReviewOrchestrator extends DurableObject { // --------------------------------------------------------------------------- // cloud-agent-next follow-up flow (session continuation) - // Uses sendMessageV2 to reuse an existing session from a previous review. + // Uses sendMessageV2Internal to reuse an existing session from a previous review. // Falls back to fresh session (prepareSession + initiate) on failure. // --------------------------------------------------------------------------- /** * Orchestration via cloud-agent-next with session continuation. - * Calls sendMessageV2 on an existing session from a previous review. + * Calls sendMessageV2Internal on an existing session from a previous review. * On failure (404, 409, etc.), falls back to runWithCloudAgentNext() for a fresh session. */ private async runWithCloudAgentNextFollowup(): Promise { @@ -1282,10 +1297,13 @@ export class CodeReviewOrchestrator extends DurableObject { } const client = this.getCloudAgentNextClient(); - console.log('[CodeReviewOrchestrator] Attempting session continuation via sendMessageV2', { - reviewId: this.state.reviewId, - previousCloudAgentSessionId: previousSessionId, - }); + console.log( + '[CodeReviewOrchestrator] Attempting session continuation via sendMessageV2Internal', + { + reviewId: this.state.reviewId, + previousCloudAgentSessionId: previousSessionId, + } + ); try { const statusUpdateResult = await this.updateStatus('running'); @@ -1339,9 +1357,6 @@ export class CodeReviewOrchestrator extends DurableObject { internalHeaders['x-skip-balance-check'] = 'true'; } - // Step 1: Update callback target via updateSession (internal-only endpoint). - // callbackTarget must be set through an internal procedure, not the - // user-facing sendMessageV2, to prevent SSRF via arbitrary callback URLs. const callbackTarget = await callbackTargetForAttempt( this.env.API_URL, this.state.reviewId, @@ -1349,26 +1364,26 @@ export class CodeReviewOrchestrator extends DurableObject { this.env.CALLBACK_TOKEN_SECRET ); - await client.updateSession(internalHeaders, { - cloudAgentSessionId: previousSessionId, - callbackTarget, - }); - - // Step 2: Send follow-up message (user-facing, no callbackTarget) - console.log('[CodeReviewOrchestrator] Calling sendMessageV2', { + console.log('[CodeReviewOrchestrator] Calling sendMessageV2Internal', { reviewId: this.state.reviewId, cloudAgentSessionId: previousSessionId, callbackUrl: callbackTarget.url, }); - const sendResult = await client.sendMessageV2(userHeaders, { + const sendResult = await client.sendMessageV2Internal(internalHeaders, { cloudAgentSessionId: previousSessionId, prompt: this.state.sessionInput.prompt, - mode: this.state.sessionInput.mode, + mode: CODE_REVIEW_RUNTIME_AGENT_SLUG, model: this.state.sessionInput.model, variant: this.state.sessionInput.variant, githubToken: this.state.sessionInput.githubToken, gitToken: this.state.sessionInput.gitToken, + completion: { + callbackTarget, + ...(this.state.sessionInput.gateThreshold !== undefined + ? { gateThreshold: this.state.sessionInput.gateThreshold } + : {}), + }, }); // Store session ID (reusing the previous one) and execution ID @@ -1376,12 +1391,15 @@ export class CodeReviewOrchestrator extends DurableObject { sessionId: previousSessionId, }); - console.log('[CodeReviewOrchestrator] Follow-up execution started via sendMessageV2', { - reviewId: this.state.reviewId, - cloudAgentSessionId: previousSessionId, - executionId: sendResult.executionId, - status: sendResult.status, - }); + console.log( + '[CodeReviewOrchestrator] Follow-up execution started via sendMessageV2Internal', + { + reviewId: this.state.reviewId, + cloudAgentSessionId: previousSessionId, + executionId: sendResult.executionId, + status: sendResult.status, + } + ); // Done — cloud-agent-next callback will deliver terminal status } catch (error) { diff --git a/services/code-review-infra/test/integration/code-review-orchestrator.test.ts b/services/code-review-infra/test/integration/code-review-orchestrator.test.ts index 23bb28d918..422c050357 100644 --- a/services/code-review-infra/test/integration/code-review-orchestrator.test.ts +++ b/services/code-review-infra/test/integration/code-review-orchestrator.test.ts @@ -87,8 +87,18 @@ function mockSuccessfulCloudAgentNextRun() { return fetchMock; } +function requestUrl(request: RequestInfo | URL): string { + return request instanceof Request ? request.url : String(request); +} + function fetchCalls(fetchMock: ReturnType, path: string) { - return fetchMock.mock.calls.filter(([request]) => String(request).includes(path)); + return fetchMock.mock.calls.filter(([request]) => { + const url = requestUrl(request); + if (path.startsWith('/trpc/')) { + return new URL(url).pathname === path; + } + return url.includes(path); + }); } function hasFetchCall(fetchMock: ReturnType, path: string): boolean { @@ -221,6 +231,36 @@ describe('CodeReviewOrchestrator recovery', () => { headers: { 'X-Callback-Token': expectedCallbackToken }, }); expect(prepareBody.callbackTarget.headers).not.toHaveProperty('X-Internal-Secret'); + expect(prepareBody).toMatchObject({ + mode: 'code-review', + managedSession: { + executionPolicy: { + name: 'code-review-read-only', + runtime: { nonInteractive: true }, + interaction: { + question: 'deny', + permission: 'auto-reject', + terminal: 'deny', + }, + tools: { + question: false, + plan_enter: false, + plan_exit: false, + }, + }, + }, + runtimeAgents: [ + { + slug: 'code-review', + name: 'Code Review', + config: { + mode: 'primary', + model: 'test-model', + permission: expect.objectContaining({ edit: 'deny', question: 'deny' }), + }, + }, + ], + }); }); it('fresh attempt dispatch does not reuse failed state from an earlier attempt', async () => { @@ -337,7 +377,7 @@ describe('CodeReviewOrchestrator recovery', () => { expect(fetchCalls(fetchMock, '/trpc/initiateFromKilocodeSessionV2')).toHaveLength(0); expect(fetchCalls(fetchMock, '/trpc/getSessionHealth')).toHaveLength(0); expect(fetchCalls(fetchMock, '/trpc/updateSession')).toHaveLength(0); - expect(fetchCalls(fetchMock, '/trpc/sendMessageV2')).toHaveLength(0); + expect(fetchCalls(fetchMock, '/trpc/sendMessageV2Internal')).toHaveLength(0); }); it('retry-fresh starts a new retry attempt durable object instead of resetting the failed attempt', async () => { @@ -932,10 +972,7 @@ describe('CodeReviewOrchestrator recovery', () => { executionHealth: 'none', }); } - if (url.includes('/trpc/updateSession')) { - return trpcSuccess({ success: true }); - } - if (url.includes('/trpc/sendMessageV2')) { + if (url.includes('/trpc/sendMessageV2Internal')) { return trpcSuccess({ executionId: 'exec-followup', status: 'running' }); } return new Response('unexpected fetch', { status: 500 }); @@ -947,6 +984,7 @@ describe('CodeReviewOrchestrator recovery', () => { 'state', codeReview({ previousCloudAgentSessionId: previousSessionId, + sessionInput: { ...sessionInput(), gateThreshold: 'warning' }, }) ); await state.storage.setAlarm(Date.now() + 30_000); @@ -962,10 +1000,26 @@ describe('CodeReviewOrchestrator recovery', () => { }); expect(status.cliSessionId).toBeUndefined(); expect(hasFetchCall(fetchMock, '/trpc/getSessionHealth')).toBe(true); - expect(hasFetchCall(fetchMock, '/trpc/updateSession')).toBe(true); - expect(hasFetchCall(fetchMock, '/trpc/sendMessageV2')).toBe(true); + expect(hasFetchCall(fetchMock, '/trpc/updateSession')).toBe(false); + expect(hasFetchCall(fetchMock, '/trpc/sendMessageV2Internal')).toBe(true); + expect(hasFetchCall(fetchMock, '/trpc/sendMessageV2')).toBe(false); expect(hasFetchCall(fetchMock, '/trpc/prepareSession')).toBe(false); expect(hasFetchCall(fetchMock, '/trpc/initiateFromKilocodeSessionV2')).toBe(false); + + const sendCall = getFetchCall(fetchMock, '/trpc/sendMessageV2Internal'); + const sendBody = JSON.parse(String(sendCall?.[1]?.body)); + expect(sendBody).toMatchObject({ + cloudAgentSessionId: previousSessionId, + mode: 'code-review', + completion: { + callbackTarget: { + url: expect.stringContaining('/api/internal/code-review-status/'), + headers: { 'X-Callback-Token': expect.stringMatching(/^[0-9a-f]{64}$/) }, + }, + gateThreshold: 'warning', + }, + }); + expect(sendBody.completion.callbackTarget.headers).not.toHaveProperty('X-Internal-Secret'); }); it('skips continuation and prepares a fresh session when previous sandbox is unreachable', async () => { @@ -1036,8 +1090,7 @@ describe('CodeReviewOrchestrator recovery', () => { executionHealth: 'none', }); } - if (url.includes('/trpc/updateSession')) return trpcSuccess({ success: true }); - if (url.includes('/trpc/sendMessageV2')) { + if (url.includes('/trpc/sendMessageV2Internal')) { return trpcSuccess({ executionId: 'msg_followup', status: 'started', @@ -1058,7 +1111,8 @@ describe('CodeReviewOrchestrator recovery', () => { await runDurableObjectAlarm(stub); - expect(hasFetchCall(fetchMock, '/trpc/sendMessageV2')).toBe(true); + expect(hasFetchCall(fetchMock, '/trpc/sendMessageV2Internal')).toBe(true); + expect(hasFetchCall(fetchMock, '/trpc/sendMessageV2')).toBe(false); expect(hasFetchCall(fetchMock, '/trpc/prepareSession')).toBe(false); }); @@ -1111,7 +1165,7 @@ describe('CodeReviewOrchestrator recovery', () => { cliSessionId: 'ses_fresh_stale', }); expect(hasFetchCall(fetchMock, '/trpc/updateSession')).toBe(false); - expect(hasFetchCall(fetchMock, '/trpc/sendMessageV2')).toBe(false); + expect(hasFetchCall(fetchMock, '/trpc/sendMessageV2Internal')).toBe(false); expect(hasFetchCall(fetchMock, '/trpc/prepareSession')).toBe(true); }); @@ -1165,7 +1219,7 @@ describe('CodeReviewOrchestrator recovery', () => { cliSessionId: 'ses_fresh_active', }); expect(hasFetchCall(fetchMock, '/trpc/updateSession')).toBe(false); - expect(hasFetchCall(fetchMock, '/trpc/sendMessageV2')).toBe(false); + expect(hasFetchCall(fetchMock, '/trpc/sendMessageV2Internal')).toBe(false); expect(hasFetchCall(fetchMock, '/trpc/prepareSession')).toBe(true); }); @@ -1214,11 +1268,11 @@ describe('CodeReviewOrchestrator recovery', () => { }); expect(hasFetchCall(fetchMock, '/trpc/getSessionHealth')).toBe(true); expect(hasFetchCall(fetchMock, '/trpc/updateSession')).toBe(false); - expect(hasFetchCall(fetchMock, '/trpc/sendMessageV2')).toBe(false); + expect(hasFetchCall(fetchMock, '/trpc/sendMessageV2Internal')).toBe(false); expect(hasFetchCall(fetchMock, '/trpc/prepareSession')).toBe(true); }); - it('falls back to a fresh session when sendMessageV2 fails after healthy preflight', async () => { + it('falls back to a fresh session when sendMessageV2Internal fails after healthy preflight', async () => { const stub = getReviewStub(); const previousSessionId = 'agent_previous_send_failure'; const fetchMock = vi.fn(async (request: RequestInfo | URL) => { @@ -1233,10 +1287,7 @@ describe('CodeReviewOrchestrator recovery', () => { executionHealth: 'none', }); } - if (url.includes('/trpc/updateSession')) { - return trpcSuccess({ success: true }); - } - if (url.includes('/trpc/sendMessageV2')) { + if (url.includes('/trpc/sendMessageV2Internal')) { return new Response('Session not found', { status: 404 }); } if (url.includes('/trpc/prepareSession')) { @@ -1272,23 +1323,26 @@ describe('CodeReviewOrchestrator recovery', () => { cliSessionId: 'ses_fresh_after_send_failure', }); expect(hasFetchCall(fetchMock, '/trpc/getSessionHealth')).toBe(true); - expect(hasFetchCall(fetchMock, '/trpc/updateSession')).toBe(true); - expect(hasFetchCall(fetchMock, '/trpc/sendMessageV2')).toBe(true); + expect(hasFetchCall(fetchMock, '/trpc/updateSession')).toBe(false); + expect(hasFetchCall(fetchMock, '/trpc/sendMessageV2Internal')).toBe(true); + expect(hasFetchCall(fetchMock, '/trpc/sendMessageV2')).toBe(false); expect(hasFetchCall(fetchMock, '/trpc/prepareSession')).toBe(true); - const updateCall = getFetchCall(fetchMock, '/trpc/updateSession'); - const updateBody = JSON.parse(String(updateCall?.[1]?.body)); - expect(updateBody).toMatchObject({ + const sendCall = getFetchCall(fetchMock, '/trpc/sendMessageV2Internal'); + const sendBody = JSON.parse(String(sendCall?.[1]?.body)); + expect(sendBody).toMatchObject({ cloudAgentSessionId: previousSessionId, - callbackTarget: { - url: expect.stringContaining('/api/internal/code-review-status/'), - headers: { 'X-Callback-Token': expect.stringMatching(/^[0-9a-f]{64}$/) }, + completion: { + callbackTarget: { + url: expect.stringContaining('/api/internal/code-review-status/'), + headers: { 'X-Callback-Token': expect.stringMatching(/^[0-9a-f]{64}$/) }, + }, }, }); - expect(updateBody.callbackTarget.headers).not.toHaveProperty('X-Internal-Secret'); + expect(sendBody.completion.callbackTarget.headers).not.toHaveProperty('X-Internal-Secret'); }); - it('retries with a fresh session when sendMessageV2 fails with a sandbox 500', async () => { + it('retries with a fresh session when sendMessageV2Internal fails with a sandbox 500', async () => { const stub = getReviewStub(); const previousSessionId = 'agent_previous_sandbox_500'; const fetchMock = vi.fn(async (request: RequestInfo | URL) => { @@ -1303,10 +1357,7 @@ describe('CodeReviewOrchestrator recovery', () => { executionHealth: 'none', }); } - if (url.includes('/trpc/updateSession')) { - return trpcSuccess({ success: true }); - } - if (url.includes('/trpc/sendMessageV2')) { + if (url.includes('/trpc/sendMessageV2Internal')) { return trpcError(500, 'Container failed with internal server error status: 500'); } if (url.includes('/trpc/prepareSession')) { @@ -1343,8 +1394,9 @@ describe('CodeReviewOrchestrator recovery', () => { cliSessionId: 'ses_fresh_after_sandbox_500', }); expect(fetchCalls(fetchMock, '/trpc/getSessionHealth')).toHaveLength(1); - expect(fetchCalls(fetchMock, '/trpc/updateSession')).toHaveLength(1); - expect(fetchCalls(fetchMock, '/trpc/sendMessageV2')).toHaveLength(1); + expect(fetchCalls(fetchMock, '/trpc/updateSession')).toHaveLength(0); + expect(fetchCalls(fetchMock, '/trpc/sendMessageV2Internal')).toHaveLength(1); + expect(fetchCalls(fetchMock, '/trpc/sendMessageV2')).toHaveLength(0); expect(fetchCalls(fetchMock, '/trpc/prepareSession')).toHaveLength(1); expect(fetchCalls(fetchMock, '/trpc/initiateFromKilocodeSessionV2')).toHaveLength(1); @@ -1357,7 +1409,7 @@ describe('CodeReviewOrchestrator recovery', () => { }); }); - it('fails with sandbox_error when sendMessageV2 retry also hits a sandbox 500', async () => { + it('fails with sandbox_error when sendMessageV2Internal retry also hits a sandbox 500', async () => { const stub = getReviewStub(); const previousSessionId = 'agent_previous_sandbox_repeat'; const fetchMock = vi.fn(async (request: RequestInfo | URL) => { @@ -1372,10 +1424,7 @@ describe('CodeReviewOrchestrator recovery', () => { executionHealth: 'none', }); } - if (url.includes('/trpc/updateSession')) { - return trpcSuccess({ success: true }); - } - if (url.includes('/trpc/sendMessageV2')) { + if (url.includes('/trpc/sendMessageV2Internal')) { return trpcError(500, 'SandboxError: HTTP error! status: 500 during resume'); } if (url.includes('/trpc/prepareSession')) { @@ -1402,7 +1451,8 @@ describe('CodeReviewOrchestrator recovery', () => { status: 'failed', terminalReason: 'sandbox_error', }); - expect(fetchCalls(fetchMock, '/trpc/sendMessageV2')).toHaveLength(1); + expect(fetchCalls(fetchMock, '/trpc/sendMessageV2Internal')).toHaveLength(1); + expect(fetchCalls(fetchMock, '/trpc/sendMessageV2')).toHaveLength(0); expect(fetchCalls(fetchMock, '/trpc/prepareSession')).toHaveLength(1); expect(fetchCalls(fetchMock, '/trpc/initiateFromKilocodeSessionV2')).toHaveLength(0); }); From 5c9415d7b0672268c0a80312363503f8ae61d70d Mon Sep 17 00:00:00 2001 From: Alex Alecu Date: Thu, 28 May 2026 09:50:28 +0300 Subject: [PATCH 4/4] fix(code-review): preserve managed continuation callbacks Reuse persisted agent defaults for follow-ups and avoid duplicating assistant response text in completion metadata. --- .../src/cloud-agent-next-client.ts | 8 ++++++- .../cloud-agent-next/src/callbacks/types.ts | 1 - .../src/persistence/CloudAgentSession.ts | 1 - .../src/router/schemas.test.ts | 21 ++++++++++++++++ .../cloud-agent-next/src/router/schemas.ts | 13 ++++++++-- .../src/session/message-settlement-outbox.ts | 16 +++++++++---- .../src/session/session-message-queue.test.ts | 24 +++++++++++++++++++ .../session/callback-notification.test.ts | 1 + .../session/message-terminalization.test.ts | 1 + .../src/code-review-orchestrator.ts | 3 --- .../code-review-orchestrator.test.ts | 7 +++++- 11 files changed, 83 insertions(+), 13 deletions(-) diff --git a/packages/worker-utils/src/cloud-agent-next-client.ts b/packages/worker-utils/src/cloud-agent-next-client.ts index 08022df837..3b3e3efefa 100644 --- a/packages/worker-utils/src/cloud-agent-next-client.ts +++ b/packages/worker-utils/src/cloud-agent-next-client.ts @@ -95,7 +95,13 @@ export type CloudAgentSendMessageInput = { gitToken?: string; }; -export type CloudAgentSendMessageInternalInput = CloudAgentSendMessageInput & { +export type CloudAgentSendMessageInternalInput = Omit< + CloudAgentSendMessageInput, + 'mode' | 'model' | 'variant' +> & { + mode?: string; + model?: string; + variant?: string; completion?: CloudAgentMessageCompletion; }; diff --git a/services/cloud-agent-next/src/callbacks/types.ts b/services/cloud-agent-next/src/callbacks/types.ts index 8f3da44b46..6e6ab6f202 100644 --- a/services/cloud-agent-next/src/callbacks/types.ts +++ b/services/cloud-agent-next/src/callbacks/types.ts @@ -23,7 +23,6 @@ export type ExecutionCallbackPayload = { lastAssistantMessageText?: string; completionData?: { assistantMessageId?: string; - lastAssistantMessageText?: string; gateResult?: 'pass' | 'fail'; lastSeenBranch?: string; completedAt?: number; diff --git a/services/cloud-agent-next/src/persistence/CloudAgentSession.ts b/services/cloud-agent-next/src/persistence/CloudAgentSession.ts index bd5fd74331..839738df91 100644 --- a/services/cloud-agent-next/src/persistence/CloudAgentSession.ts +++ b/services/cloud-agent-next/src/persistence/CloudAgentSession.ts @@ -312,7 +312,6 @@ export class CloudAgentSession extends DurableObject { status === 'completed' ? { ...(gateResult !== undefined ? { gateResult } : {}), - ...(lastAssistantMessageText !== undefined ? { lastAssistantMessageText } : {}), ...(metadata.repository?.upstreamBranch !== undefined ? { lastSeenBranch: metadata.repository.upstreamBranch } : {}), diff --git a/services/cloud-agent-next/src/router/schemas.test.ts b/services/cloud-agent-next/src/router/schemas.test.ts index eccf0fdd94..72a0847220 100644 --- a/services/cloud-agent-next/src/router/schemas.test.ts +++ b/services/cloud-agent-next/src/router/schemas.test.ts @@ -6,6 +6,7 @@ import { LegacyExecutionResponse, SendMessageInput, SendMessageV2Input, + TrustedSendMessageV2Input, StartSessionOutput, StartSessionInput, } from './schemas.js'; @@ -157,6 +158,26 @@ describe('sendMessageV2 input compatibility', () => { images: validImages, }); }); + + it('allows trusted continuations to reuse the stored agent selection', () => { + const result = TrustedSendMessageV2Input.safeParse({ + cloudAgentSessionId: validSessionId, + prompt: 'continue managed review', + completion: { gateThreshold: 'warning' }, + }); + + expect(result.success).toBe(true); + if (!result.success) return; + + expect(result.data).toMatchObject({ + cloudAgentSessionId: validSessionId, + prompt: 'continue managed review', + completion: { gateThreshold: 'warning' }, + }); + expect(result.data).not.toHaveProperty('mode'); + expect(result.data).not.toHaveProperty('model'); + expect(result.data).not.toHaveProperty('variant'); + }); }); describe('message ID schema validation', () => { diff --git a/services/cloud-agent-next/src/router/schemas.ts b/services/cloud-agent-next/src/router/schemas.ts index c3efa1ac24..4dec178024 100644 --- a/services/cloud-agent-next/src/router/schemas.ts +++ b/services/cloud-agent-next/src/router/schemas.ts @@ -237,9 +237,18 @@ const SendMessageV2PromptPayloadInput = SendMessageV2Options.extend({ const SendMessageV2CommandPayloadInput = SendMessageV2Options.extend({ payload: CommandSendPayload, }); -const TrustedSendMessageV2FlatInput = TrustedSendMessageV2Options.extend(PromptPayload.shape); +const TrustedSendMessageV2FlatInput = TrustedSendMessageV2Options.extend({ + prompt: PromptPayload.shape.prompt, + mode: ModeSlugSchema.optional(), + model: modelIdSchema.optional(), + variant: PromptPayload.shape.variant, +}); +const TrustedPromptSendPayload = PromptSendPayload.extend({ + mode: ModeSlugSchema.optional(), + model: modelIdSchema.optional(), +}); const TrustedSendMessageV2PromptPayloadInput = TrustedSendMessageV2Options.extend({ - payload: PromptSendPayload, + payload: TrustedPromptSendPayload, }); const TrustedSendMessageV2CommandPayloadInput = TrustedSendMessageV2Options.extend({ payload: CommandSendPayload, diff --git a/services/cloud-agent-next/src/session/message-settlement-outbox.ts b/services/cloud-agent-next/src/session/message-settlement-outbox.ts index 169ae3512c..afc62cf2a9 100644 --- a/services/cloud-agent-next/src/session/message-settlement-outbox.ts +++ b/services/cloud-agent-next/src/session/message-settlement-outbox.ts @@ -104,6 +104,15 @@ function extractAssistantTextFromParts(parts: AssistantMessagePart[]): string { return pieces.join('').trim(); } +function withoutDuplicatedAssistantText( + completionData: SessionMessageState['completionData'] +): SessionMessageState['completionData'] { + if (completionData?.lastAssistantMessageText === undefined) return completionData; + const sanitized = { ...completionData }; + delete sanitized.lastAssistantMessageText; + return sanitized; +} + function redactCallbackTargetUrl(callbackUrl: string): string { try { const url = new URL(callbackUrl); @@ -275,7 +284,7 @@ export function createMessageSettlementOutbox( completionSource: state.completionSource, }; const completionData = { - ...state.completionData, + ...withoutDuplicatedAssistantText(state.completionData), ...(extra?.gateResult !== undefined ? { gateResult: extra.gateResult } : {}), }; if (state.assistantMessageId) { @@ -382,14 +391,13 @@ export function createMessageSettlementOutbox( const completionData = status === 'completed' ? { - ...state.completionData, - ...(lastAssistantMessageText !== undefined ? { lastAssistantMessageText } : {}), + ...withoutDuplicatedAssistantText(state.completionData), ...(metadata?.repository?.upstreamBranch !== undefined ? { lastSeenBranch: metadata.repository.upstreamBranch } : {}), ...(state.gateResult !== undefined ? { gateResult: state.gateResult } : {}), } - : state.completionData; + : withoutDuplicatedAssistantText(state.completionData); if (status === 'completed' && completionData !== state.completionData) { state.completionData = completionData; diff --git a/services/cloud-agent-next/src/session/session-message-queue.test.ts b/services/cloud-agent-next/src/session/session-message-queue.test.ts index 6db39109a0..4de6227b85 100644 --- a/services/cloud-agent-next/src/session/session-message-queue.test.ts +++ b/services/cloud-agent-next/src/session/session-message-queue.test.ts @@ -550,6 +550,30 @@ describe('SessionMessageQueue', () => { await expect(listPendingSessionMessages(harness.storage)).resolves.toHaveLength(0); }); + it('uses stored agent selection for managed submissions without overrides', async () => { + const harness = createQueueHarness({ + metadata: createMetadata({ + agent: { mode: 'code-review', model: 'default-model', variant: 'alpha' }, + executionPolicy: ManagedSessionExecutionPolicySchema.parse( + buildCodeReviewManagedSessionPolicy() + ), + }), + }); + + const result = await harness.queue.admitSubmittedMessage({ + userId: 'user_test' as UserId, + turn: { type: 'prompt', id: FIRST_MESSAGE_ID, prompt: 'continue review' }, + }); + const state = await getSessionMessageState(harness.storage, FIRST_MESSAGE_ID); + + expect(result).toMatchObject({ success: true, messageId: FIRST_MESSAGE_ID }); + expect(state?.admissionSnapshot?.agent).toEqual({ + mode: 'code-review', + model: 'default-model', + variant: 'alpha', + }); + }); + it('accepts replay matching predecessor partial immutable constraints without a stored turn', async () => { const harness = createQueueHarness(); await harness.storage.put(`session_message:${FIRST_MESSAGE_ID}`, { diff --git a/services/cloud-agent-next/test/integration/session/callback-notification.test.ts b/services/cloud-agent-next/test/integration/session/callback-notification.test.ts index c6ee61a46e..5661a016e2 100644 --- a/services/cloud-agent-next/test/integration/session/callback-notification.test.ts +++ b/services/cloud-agent-next/test/integration/session/callback-notification.test.ts @@ -140,6 +140,7 @@ describe('Callback notification with latest assistant message', () => { const [job] = queue.captured; expect(job.payload.status).toBe('completed'); expect(job.payload.lastAssistantMessageText).toBe('Hello world'); + expect(job.payload.completionData).not.toHaveProperty('lastAssistantMessageText'); expect(job.target.url).toBe('https://example.com/callback'); }); diff --git a/services/cloud-agent-next/test/integration/session/message-terminalization.test.ts b/services/cloud-agent-next/test/integration/session/message-terminalization.test.ts index 0bd87e7f1c..b1b9c99bd2 100644 --- a/services/cloud-agent-next/test/integration/session/message-terminalization.test.ts +++ b/services/cloud-agent-next/test/integration/session/message-terminalization.test.ts @@ -783,6 +783,7 @@ describe('message terminalization and stream events', () => { expect(job.payload.status).toBe('completed'); expect(job.payload.messageId).toBe('msg_018f1e2d3c4btermcorrabcd01'); expect(job.payload.lastAssistantMessageText).toBe('Correct answer'); + expect(job.payload.completionData).not.toHaveProperty('lastAssistantMessageText'); }); it('completed callback omits assistant text when no matching parentID reply exists', async () => { diff --git a/services/code-review-infra/src/code-review-orchestrator.ts b/services/code-review-infra/src/code-review-orchestrator.ts index 48350cdc83..f68bf3959b 100644 --- a/services/code-review-infra/src/code-review-orchestrator.ts +++ b/services/code-review-infra/src/code-review-orchestrator.ts @@ -1373,9 +1373,6 @@ export class CodeReviewOrchestrator extends DurableObject { const sendResult = await client.sendMessageV2Internal(internalHeaders, { cloudAgentSessionId: previousSessionId, prompt: this.state.sessionInput.prompt, - mode: CODE_REVIEW_RUNTIME_AGENT_SLUG, - model: this.state.sessionInput.model, - variant: this.state.sessionInput.variant, githubToken: this.state.sessionInput.githubToken, gitToken: this.state.sessionInput.gitToken, completion: { diff --git a/services/code-review-infra/test/integration/code-review-orchestrator.test.ts b/services/code-review-infra/test/integration/code-review-orchestrator.test.ts index 422c050357..abe27eac0b 100644 --- a/services/code-review-infra/test/integration/code-review-orchestrator.test.ts +++ b/services/code-review-infra/test/integration/code-review-orchestrator.test.ts @@ -1010,7 +1010,6 @@ describe('CodeReviewOrchestrator recovery', () => { const sendBody = JSON.parse(String(sendCall?.[1]?.body)); expect(sendBody).toMatchObject({ cloudAgentSessionId: previousSessionId, - mode: 'code-review', completion: { callbackTarget: { url: expect.stringContaining('/api/internal/code-review-status/'), @@ -1019,6 +1018,9 @@ describe('CodeReviewOrchestrator recovery', () => { gateThreshold: 'warning', }, }); + expect(sendBody).not.toHaveProperty('mode'); + expect(sendBody).not.toHaveProperty('model'); + expect(sendBody).not.toHaveProperty('variant'); expect(sendBody.completion.callbackTarget.headers).not.toHaveProperty('X-Internal-Secret'); }); @@ -1339,6 +1341,9 @@ describe('CodeReviewOrchestrator recovery', () => { }, }, }); + expect(sendBody).not.toHaveProperty('mode'); + expect(sendBody).not.toHaveProperty('model'); + expect(sendBody).not.toHaveProperty('variant'); expect(sendBody.completion.callbackTarget.headers).not.toHaveProperty('X-Internal-Secret'); });