Honor session context windows and cancel stalled MCP calls

src/commands/set-context-window/set-context-window.test.ts

@@ -0,0 +1,65 @@
+import { afterEach, beforeEach, expect, test } from 'bun:test'
+import {
+  acquireSharedMutationLock,
+  releaseSharedMutationLock,
+} from '../../test/sharedMutationLock.js'
+import { call as clearContextWindow } from '../clear-context-window/clear-context-window.js'
+import { call as setContextWindow } from './set-context-window.js'
+import {
+  clearSessionContextWindowOverride,
+  getContextWindowForModel,
+} from '../../utils/context.js'
+import {
+  getAutoCompactThreshold,
+  getEffectiveContextWindowSize,
+} from '../../services/compact/autoCompact.js'
+
+let hasSharedMutationLock = false
+const savedEnv = {
+  CLAUDE_CODE_AUTO_COMPACT_WINDOW: process.env.CLAUDE_CODE_AUTO_COMPACT_WINDOW,
+  CLAUDE_CODE_MAX_OUTPUT_TOKENS: process.env.CLAUDE_CODE_MAX_OUTPUT_TOKENS,
+  CLAUDE_AUTOCOMPACT_PCT_OVERRIDE:
+    process.env.CLAUDE_AUTOCOMPACT_PCT_OVERRIDE,
+}
+
+beforeEach(async () => {
+  await acquireSharedMutationLock('commands/set-context-window/set-context-window.test.ts')
+  hasSharedMutationLock = true
+  process.env.CLAUDE_CODE_AUTO_COMPACT_WINDOW = '100000'
+  process.env.CLAUDE_CODE_MAX_OUTPUT_TOKENS = '20000'
+  delete process.env.CLAUDE_AUTOCOMPACT_PCT_OVERRIDE
+  clearSessionContextWindowOverride('claude-sonnet-4')
+})
+
+afterEach(() => {
+  clearSessionContextWindowOverride('claude-sonnet-4')
+  for (const [key, value] of Object.entries(savedEnv)) {
+    if (value === undefined) {
+      delete process.env[key]
+    } else {
+      process.env[key] = value
+    }
+  }
+  if (hasSharedMutationLock) {
+    releaseSharedMutationLock()
+    hasSharedMutationLock = false
+  }
+})
+
+test('/set-context-window updates tokens and compaction immediately; clear restores defaults', async () => {
+  const context = {
+    options: { mainLoopModel: 'claude-sonnet-4' },
+  } as never
+
+  expect(getAutoCompactThreshold('claude-sonnet-4')).toBe(50_000)
+  const result = await setContextWindow('1000000', context)
+  expect(result).toMatchObject({ type: 'text', value: expect.stringContaining('1,000,000') })
+  expect(getContextWindowForModel('claude-sonnet-4')).toBe(1_000_000)
+  expect(getEffectiveContextWindowSize('claude-sonnet-4')).toBe(980_000)
+  expect(getAutoCompactThreshold('claude-sonnet-4')).toBe(950_000)
+
+  await clearContextWindow('claude-sonnet-4', context)
+  expect(getContextWindowForModel('claude-sonnet-4')).toBe(200_000)
+  expect(getEffectiveContextWindowSize('claude-sonnet-4')).toBe(80_000)
+  expect(getAutoCompactThreshold('claude-sonnet-4')).toBe(50_000)
+})

src/services/compact/autoCompact.test.ts

@@ -122,6 +122,7 @@ beforeEach(async () => {
     delete process.env.CLAUDE_CODE_MAX_CONTEXT_TOKENS
     delete process.env.CLAUDE_CODE_AUTO_COMPACT_WINDOW
     delete process.env.CLAUDE_CODE_MAX_OUTPUT_TOKENS
+    delete process.env.USER_TYPE
   } catch (error) {
     releaseSharedMutationLock()
     hasSharedMutationLock = false
@@ -335,6 +336,41 @@ describe('getAutoCompactThreshold', () => {
     expect(getAutoCompactThreshold('claude-sonnet-4')).toBe(20_000)
   })
 
+  test('internal context cap keeps precedence over session override and auto-compact cap', async () => {
+    process.env.USER_TYPE = 'ant'
+    process.env.CLAUDE_CODE_MAX_CONTEXT_TOKENS = '200000'
+    process.env.CLAUDE_CODE_AUTO_COMPACT_WINDOW = '100000'
+    realContext.setSessionContextWindowOverride('claude-sonnet-4', 1_000_000)
+    const { getEffectiveContextWindowSize, getAutoCompactThreshold } =
+      await importAutoCompact()
+
+    try {
+      expect(getEffectiveContextWindowSize('claude-sonnet-4')).toBe(80_000)
+      expect(getAutoCompactThreshold('claude-sonnet-4')).toBe(50_000)
+    } finally {
+      realContext.clearSessionContextWindowOverride('claude-sonnet-4')
+      delete process.env.CLAUDE_CODE_MAX_CONTEXT_TOKENS
+    }
+  })
+
+  test('session context-window override immediately updates auto-compact threshold', async () => {
+    process.env.CLAUDE_CODE_AUTO_COMPACT_WINDOW = '100000'
+    process.env.CLAUDE_CODE_MAX_OUTPUT_TOKENS = '20000'
+    delete process.env.CLAUDE_AUTOCOMPACT_PCT_OVERRIDE
+    const { getAutoCompactThreshold, getEffectiveContextWindowSize } =
+      await importAutoCompact()
+
+    expect(getAutoCompactThreshold('claude-sonnet-4')).toBe(50_000)
+    realContext.setSessionContextWindowOverride('claude-sonnet-4', 1_000_000)
+
+    try {
+      expect(getEffectiveContextWindowSize('claude-sonnet-4')).toBe(980_000)
+      expect(getAutoCompactThreshold('claude-sonnet-4')).toBe(950_000)
+    } finally {
+      realContext.clearSessionContextWindowOverride('claude-sonnet-4')
+    }
+  })
+
   test('keeps compaction and warning thresholds usable across mid-sized windows', async () => {
     process.env.CLAUDE_CODE_AUTO_COMPACT_WINDOW = '64000'
     const { calculateTokenWarningState, getAutoCompactThreshold } =

src/services/compact/autoCompact.ts

@@ -5,7 +5,10 @@ import type { QuerySource } from '../../constants/querySource.js'
 import type { ToolUseContext } from '../../Tool.js'
 import type { Message } from '../../types/message.js'
 import { getGlobalConfig } from '../../utils/config.js'
-import { getContextWindowForModel } from '../../utils/context.js'
+import {
+  getContextWindowForModel,
+  getSessionContextWindowOverride,
+} from '../../utils/context.js'
 import { logForDebugging } from '../../utils/debug.js'
 import { isEnvTruthy } from '../../utils/envUtils.js'
 import { hasExactErrorMessage } from '../../utils/errors.js'
@@ -49,8 +52,21 @@ export function getEffectiveContextWindowSize(
     runtimeLimits,
   )
 
+  // An explicit per-model session setting is more specific than the legacy
+  // auto-compact cap. Otherwise that cap can silently keep a newly configured
+  // large context window stuck at (for example) 100k for compaction decisions.
+  const internalContextWindowOverride =
+    process.env.USER_TYPE === 'ant'
+      ? parseInt(process.env.CLAUDE_CODE_MAX_CONTEXT_TOKENS ?? '', 10)
+      : NaN
+  const internalOverrideTakesPrecedence =
+    Number.isFinite(internalContextWindowOverride) &&
+    internalContextWindowOverride > 0
+  const hasSessionOverride =
+    getSessionContextWindowOverride(model) !== undefined &&
+    !internalOverrideTakesPrecedence
   const autoCompactWindow = process.env.CLAUDE_CODE_AUTO_COMPACT_WINDOW
-  if (autoCompactWindow) {
+  if (autoCompactWindow && !hasSessionOverride) {
     const parsed = parseInt(autoCompactWindow, 10)
     if (!isNaN(parsed) && parsed > 0) {
       contextWindow = Math.min(contextWindow, parsed)

src/services/mcp/client.activity.test.ts

@@ -3,15 +3,41 @@ import {
   ErrorCode,
   McpError,
 } from '@modelcontextprotocol/sdk/types.js'
+import { AbortError } from '../../utils/errors.js'
 import { createAssistantMessage } from '../../utils/messages.js'
-import { fetchToolsForClient } from './client.js'
+import {
+  callMCPToolWithUrlElicitationRetry,
+  fetchToolsForClient,
+} from './client.js'
 import type { ConnectedMCPServer } from './types.js'
 
 describe('MCP tool activity', () => {
   afterEach(() => {
     vi.useRealTimers()
   })
 
+  test('URL elicitation cancellation uses the shared AbortError type', async () => {
+    const abortController = new AbortController()
+    abortController.abort()
+    const callToolFn = vi.fn(async () => ({ content: [] }))
+
+    await expect(
+      callMCPToolWithUrlElicitationRetry({
+        client: {} as never,
+        clientConnection: {
+          type: 'connected',
+          name: 'elicitation-cancel-test',
+        } as never,
+        tool: 'slow-tool',
+        args: {},
+        signal: abortController.signal,
+        setAppState: vi.fn(),
+        callToolFn,
+      }),
+    ).rejects.toBeInstanceOf(AbortError)
+    expect(callToolFn).not.toHaveBeenCalled()
+  })
+
   test('emits progress heartbeats while a tool call remains pending', async () => {
     vi.useFakeTimers()
     let resolveToolCall: ((result: unknown) => void) | undefined
@@ -155,6 +181,79 @@ describe('MCP tool activity', () => {
     expect(onProgress).toHaveBeenCalledTimes(progressCountAfterCompletion)
   })
 
+  test('aborts the MCP request when the tool-call timeout expires', async () => {
+    vi.useFakeTimers()
+    const originalTimeout = process.env.MCP_TOOL_TIMEOUT
+    process.env.MCP_TOOL_TIMEOUT = '1000'
+    let requestSignal: AbortSignal | undefined
+    const sdkClient = {
+      request: vi.fn(async () => ({
+        tools: [
+          {
+            name: 'slow-tool',
+            inputSchema: { type: 'object' },
+          },
+        ],
+      })),
+      callTool: vi.fn(
+        (
+          _request: unknown,
+          _schema: unknown,
+          options: { signal?: AbortSignal },
+        ) => {
+          requestSignal = options.signal
+          return new Promise(() => {})
+        },
+      ),
+    }
+    const connection = {
+      type: 'connected',
+      name: 'timeout-abort-test',
+      config: { type: 'sdk' },
+      capabilities: { tools: {} },
+      client: sdkClient,
+    } as unknown as ConnectedMCPServer
+
+    try {
+      const [tool] = await fetchToolsForClient(connection)
+      expect(tool).toBeDefined()
+      const parentMessage = createAssistantMessage({
+        content: [
+          {
+            type: 'tool_use',
+            id: 'toolu_timeout_abort',
+            name: tool!.name,
+            input: {},
+          },
+        ],
+      })
+      const callPromise = tool!.call(
+        {},
+        {
+          abortController: new AbortController(),
+          setAppState: vi.fn(),
+        } as never,
+        undefined as never,
+        parentMessage,
+      )
+
+      await Promise.resolve()
+      await Promise.resolve()
+      expect(requestSignal).toBeDefined()
+      expect(requestSignal?.aborted).toBe(false)
+
+      vi.advanceTimersByTime(1000)
+      await expect(callPromise).rejects.toThrow('timed out after 1s')
+      expect(requestSignal?.aborted).toBe(true)
+    } finally {
+      if (originalTimeout === undefined) {
+        delete process.env.MCP_TOOL_TIMEOUT
+      } else {
+        process.env.MCP_TOOL_TIMEOUT = originalTimeout
+      }
+    }
+  })
+
   test('a throwing started callback does not prevent the tool call', async () => {
     const sdkClient = {
       request: vi.fn(async () => ({

src/services/mcp/client.test.ts

@@ -1,16 +1,20 @@
 import assert from 'node:assert/strict'
 import test from 'node:test'
+import { Client } from '@modelcontextprotocol/sdk/client/index.js'
 import { UnauthorizedError } from '@modelcontextprotocol/sdk/client/auth.js'
 import { SSEClientTransport } from '@modelcontextprotocol/sdk/client/sse.js'
 import { StreamableHTTPClientTransport } from '@modelcontextprotocol/sdk/client/streamableHttp.js'
+import { CallToolResultSchema } from '@modelcontextprotocol/sdk/types.js'
 
 import {
   appendBoundedMcpStderr,
+  attachMcpRequestCancellationHandler,
   buildMcpSseEventSourceHeaders,
   buildMcpSseRequestHeaders,
   cleanupFailedConnection,
   buildMcpStdioCommand,
   logMcpServerStderr,
+  wrapFetchWithRequestCancellation,
 } from './client.js'
 import { wrapFetchWithStepUpDetection } from './auth.js'
 import {
@@ -206,6 +210,79 @@ test('SSE failed recovery reaches UnauthorizedError without leaking the resource
   assert.equal(fixture.getRedirectCalls(), 1)
 })
 
+test('timed-out Streamable HTTP tool calls cancel the response stream after headers', async () => {
+  const activeRequests = new Map<string | number, AbortController>()
+  let streamCancelled = false
+  const server = Bun.serve({
+    port: 0,
+    async fetch(request) {
+      if (request.method === 'GET') {
+        return new Response(null, { status: 405 })
+      }
+      const text = await request.text()
+      if (!text) return new Response(null, { status: 202 })
+      const payload = JSON.parse(text) as { id: number; method: string }
+      if (payload.method === 'initialize') {
+        return Response.json({
+          jsonrpc: '2.0',
+          id: payload.id,
+          result: {
+            protocolVersion: '2025-03-26',
+            capabilities: {},
+            serverInfo: { name: 'timeout-stream-test', version: '1.0.0' },
+          },
+        })
+      }
+      if (payload.method === 'notifications/initialized') {
+        return new Response(null, { status: 202 })
+      }
+      if (payload.method === 'tools/call') {
+        return new Response(
+          new ReadableStream<Uint8Array>({
+            start() {},
+            cancel() {
+              streamCancelled = true
+            },
+          }),
+          {
+            headers: { 'content-type': 'text/event-stream' },
+          },
+        )
+      }
+      return new Response(null, { status: 202 })
+    },
+  })
+  const transport = new StreamableHTTPClientTransport(
+    new URL(`http://127.0.0.1:${server.port}/mcp`),
+    {
+      fetch: wrapFetchWithRequestCancellation(fetch, activeRequests),
+    },
+  )
+  attachMcpRequestCancellationHandler(transport, activeRequests)
+  const client = new Client({ name: 'timeout-test', version: '1.0.0' })
+
+  try {
+    await client.connect(transport)
+    await assert.rejects(
+      client.callTool(
+        { name: 'slow-tool', arguments: {} },
+        CallToolResultSchema,
+        { timeout: 20 },
+      ),
+    )
+
+    const cancellationDeadline = Date.now() + 1_000
+    while (!streamCancelled && Date.now() < cancellationDeadline) {
+      await new Promise(resolve => setTimeout(resolve, 5))
+    }
+    assert.equal(streamCancelled, true)
+    assert.equal(activeRequests.size, 0)
+  } finally {
+    await client.close()
+    server.stop(true)
+  }
+})
+
 test('cleanupFailedConnection awaits transport close before resolving', async () => {
   let closed = false
   let resolveClose: (() => void) | undefined