Skip to content
Closed
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
Original file line number Diff line number Diff line change
Expand Up @@ -121,6 +121,7 @@ const startupDependencies = Layer.mergeAll(
getInstanceInfo: () => Effect.die("unused"),
rollbackConversation: () => Effect.die("unused"),
uploadFeedback: () => Effect.die("unused"),
getContextUsage: () => Effect.die("unused"),
streamEvents: Stream.empty,
}),
);
Expand Down
1 change: 1 addition & 0 deletions apps/server/src/auth/RpcAuthorization.ts
Original file line number Diff line number Diff line change
Expand Up @@ -106,6 +106,7 @@ export const RPC_REQUIRED_SCOPES = {
[WS_METHODS.attachmentsCreateUploadUrl]: AuthOrchestrationOperateScope,
[WS_METHODS.attachmentsDelete]: AuthOrchestrationOperateScope,
[WS_METHODS.providerUploadFeedback]: AuthOrchestrationOperateScope,
[WS_METHODS.providerGetContextUsage]: AuthOrchestrationReadScope,
[WS_METHODS.subscribeVcsStatus]: AuthOrchestrationReadScope,
[WS_METHODS.subscribeResourceTelemetry]: AuthOrchestrationReadScope,
[WS_METHODS.vcsRefreshStatus]: AuthOrchestrationReadScope,
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -138,6 +138,7 @@ function createProviderServiceHarness(
}),
rollbackConversation,
uploadFeedback: () => unsupported(),
getContextUsage: () => unsupported(),
get streamEvents() {
return Stream.fromPubSub(runtimeEventPubSub);
},
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -391,6 +391,7 @@ describe("ProviderCommandReactor", () => {
},
rollbackConversation: () => unsupported(),
uploadFeedback: () => unsupported(),
getContextUsage: () => unsupported(),
get streamEvents() {
return Stream.fromPubSub(runtimeEventPubSub);
},
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -142,6 +142,7 @@ function createProviderServiceHarness() {
},
rollbackConversation: () => unsupported(),
uploadFeedback: () => unsupported(),
getContextUsage: () => unsupported(),
get streamEvents() {
return Stream.fromPubSub(runtimeEventPubSub).pipe(
Stream.flatMap(({ events, enqueued }) =>
Expand Down
103 changes: 93 additions & 10 deletions apps/server/src/provider/Layers/ClaudeAdapter.test.ts
Original file line number Diff line number Diff line change
Expand Up @@ -8,6 +8,7 @@ import type {
Options as ClaudeQueryOptions,
PermissionMode,
PermissionResult,
SDKControlGetContextUsageResponse,
SDKMessage,
SDKUserMessage,
} from "@anthropic-ai/claude-agent-sdk";
Expand Down Expand Up @@ -118,6 +119,10 @@ class FakeClaudeQuery implements AsyncIterable<SDKMessage> {
this.setMaxThinkingTokensCalls.push(maxThinkingTokens);
};

getContextUsage = async (): Promise<SDKControlGetContextUsageResponse> => {
throw new Error("getContextUsage not stubbed");
};

readonly close = (): void => {
this.closeCalls += 1;
if (this.closeError !== undefined) {
Expand Down Expand Up @@ -2926,16 +2931,14 @@ describe("ClaudeAdapterLive", () => {
it.effect("completes with result usage without querying current context usage", () => {
const harness = makeHarness();
let getContextUsageCalls = 0;
Object.assign(harness.query, {
getContextUsage: async () => {
getContextUsageCalls += 1;
return {
totalTokens: 999,
maxTokens: 200000,
isAutoCompactEnabled: true,
};
},
});
harness.query.getContextUsage = async () => {
getContextUsageCalls += 1;
return {
totalTokens: 999,
maxTokens: 200000,
isAutoCompactEnabled: true,
} as unknown as SDKControlGetContextUsageResponse;
};
return Effect.gen(function* () {
const adapter = yield* ClaudeAdapter;
const runtimeEventsFiber = yield* Stream.take(adapter.streamEvents, 7).pipe(
Expand Down Expand Up @@ -3238,6 +3241,86 @@ describe("ClaudeAdapterLive", () => {
);
});

it.effect("reads the context breakdown on demand through the SDK", () => {
const harness = makeHarness();
let getContextUsageCalls = 0;
harness.query.getContextUsage = async () => {
getContextUsageCalls += 1;
return {
categories: [
{ name: "System tools", tokens: 19_200, color: "inactive" },
{ name: "Messages", tokens: 12_000.4, color: "purple" },
{ name: "Free space", tokens: 168_800, color: "inactive" },
{ name: "MCP tools (deferred)", tokens: 42_400, color: "inactive", isDeferred: true },
],
totalTokens: 31_200,
maxTokens: 200_000,
rawMaxTokens: 200_000,
percentage: 15.6,
gridRows: [],
model: "claude-fable-5-1",
memoryFiles: [{ path: "/Users/me/repo/CLAUDE.md", type: "Project", tokens: 3_100 }],
mcpTools: [
{ name: "mcp__t3-code__preview_click", serverName: "t3-code", tokens: 400 },
{ name: "mcp__xcodebuildmcp__build_run_sim", serverName: "xcodebuildmcp", tokens: 900 },
],
agents: [{ agentType: "sol", source: "user", tokens: 87 }],
systemTools: [{ name: "Bash", tokens: 5_000 }],
isAutoCompactEnabled: true,
apiUsage: null,
} as SDKControlGetContextUsageResponse;
};
return Effect.gen(function* () {
const adapter = yield* ClaudeAdapter;
yield* adapter.startSession({
threadId: THREAD_ID,
provider: ProviderDriverKind.make("claudeAgent"),
runtimeMode: "full-access",
});

// Nothing queried the breakdown yet: it is on-demand only.
assert.equal(getContextUsageCalls, 0);
assert.isDefined(adapter.getContextUsage);
const usage = yield* adapter.getContextUsage!(THREAD_ID);

assert.equal(getContextUsageCalls, 1);
assert.equal(usage.model, "claude-fable-5-1");
assert.equal(usage.totalTokens, 31_200);
assert.equal(usage.maxTokens, 200_000);
assert.deepEqual(usage.categories, [
{ name: "System tools", tokens: 19_200, deferred: false },
{ name: "Messages", tokens: 12_000, deferred: false },
{ name: "Free space", tokens: 168_800, deferred: false },
{ name: "MCP tools (deferred)", tokens: 42_400, deferred: true },
]);
assert.deepEqual(
usage.groups.map((group) => [group.label, group.tokens, group.items.length]),
[
["System tools", 5_000, 1],
["MCP tools", 1_300, 2],
["Memory files", 3_100, 1],
["Custom agents", 87, 1],
],
);
// Items inside a group sort by size so the biggest offender is first,
// and MCP names drop their `mcp__<server>__` prefix.
assert.deepEqual(usage.groups[1]?.items[0], {
name: "build_run_sim",
tokens: 900,
detail: "xcodebuildmcp",
});
// Memory file paths shrink to their last two segments.
assert.deepEqual(usage.groups[2]?.items[0], {
name: "repo/CLAUDE.md",
tokens: 3_100,
detail: "Project",
});
}).pipe(
Effect.provideService(Random.Random, makeDeterministicRandomService()),
Effect.provide(harness.layer),
);
});

it.effect("preserves compacted usage when completion follows an older assistant frame", () => {
const harness = makeHarness();
return Effect.gen(function* () {
Expand Down
118 changes: 118 additions & 0 deletions apps/server/src/provider/Layers/ClaudeAdapter.ts
Original file line number Diff line number Diff line change
Expand Up @@ -13,6 +13,7 @@ import {
type PermissionMode,
type PermissionResult,
type PermissionUpdate,
type SDKControlGetContextUsageResponse,
type SDKMessage,
type SDKRateLimitInfo,
type SDKResultMessage,
Expand All @@ -31,6 +32,7 @@ import {
type ClaudeSettings,
EventId,
type ProviderApprovalDecision,
type ProviderContextUsage,
ProviderDriverKind,
ProviderInstanceId,
type ModelSelection,
Expand Down Expand Up @@ -339,9 +341,111 @@ interface ClaudeQueryRuntime extends AsyncIterable<SDKMessage> {
readonly setModel: (model?: string) => Promise<void>;
readonly setPermissionMode: (mode: PermissionMode) => Promise<void>;
readonly setMaxThinkingTokens: (maxThinkingTokens: number | null) => Promise<void>;
readonly getContextUsage: () => Promise<SDKControlGetContextUsageResponse>;
readonly close: () => void;
}

/**
* Translate the SDK's `/context` report into the wire contract. The SDK's
* category list is ordered by size and already includes free space and any
* deferred slices; the detail groups mirror the expandable rows Claude Code
* shows under that list.
*/
export function toProviderContextUsage(
response: SDKControlGetContextUsageResponse,
): ProviderContextUsage {
const categories = response.categories
.filter((category) => category.name.trim().length > 0 && category.tokens >= 0)
.map((category) => ({
name: category.name,
tokens: Math.round(category.tokens),
deferred: category.isDeferred === true,
}));

const groups: Array<ProviderContextUsage["groups"][number]> = [];
const pushGroup = (
label: string,
items: ReadonlyArray<{ name: string; tokens: number; detail?: string }>,
) => {
const kept = items
.filter((item) => item.name.trim().length > 0 && item.tokens >= 0)
.map((item) => ({
name: item.name,
tokens: Math.round(item.tokens),
...(item.detail && item.detail.trim().length > 0 ? { detail: item.detail } : {}),
}))
.sort((left, right) => right.tokens - left.tokens);
if (kept.length === 0) {
return;
}
groups.push({
label,
tokens: kept.reduce((sum, item) => sum + item.tokens, 0),
items: kept,
});
};

pushGroup(
"System prompt",
(response.systemPromptSections ?? []).map((section) => ({
name: section.name,
tokens: section.tokens,
})),
);
pushGroup(
"System tools",
(response.systemTools ?? []).map((tool) => ({ name: tool.name, tokens: tool.tokens })),
);
// MCP tool names carry an `mcp__<server>__` prefix; the server already
// appears as the detail column, so only the bare tool name is shown.
pushGroup(
"MCP tools",
response.mcpTools.map((tool) => ({
name: tool.name.replace(/^mcp__[^_]+(?:_[^_]+)*__/, "") || tool.name,
tokens: tool.tokens,
detail: tool.serverName,
})),
);
// Memory files arrive as absolute paths; the last two segments identify
// them (`t3code/AGENTS.md`) without the home-directory prefix that would
// otherwise get truncated away in a narrow popover.
pushGroup(
"Memory files",
response.memoryFiles.map((file) => ({
name: file.path.split(/[\\/]/).filter(Boolean).slice(-2).join("/") || file.path,
tokens: file.tokens,
detail: file.type,
})),
);
pushGroup(
"Skills",
(response.skills?.skillFrontmatter ?? []).map((skill) => ({
name: skill.name,
tokens: skill.tokens,
detail: skill.source,
})),
);
pushGroup(
"Custom agents",
response.agents.map((agent) => ({
name: agent.agentType,
tokens: agent.tokens,
detail: agent.source,
})),
);
// `messageBreakdown` is left out on purpose: it totals every API call in
// the session, so it exceeds the resident "Messages" slice and would read
// as a contradiction next to it.

return {
model: response.model,
totalTokens: Math.max(0, Math.round(response.totalTokens)),
maxTokens: Math.max(1, Math.round(response.maxTokens)),
categories,
groups,
};
}

export interface ClaudeAdapterLiveOptions {
readonly instanceId?: ProviderInstanceId;
readonly environment?: NodeJS.ProcessEnv;
Expand Down Expand Up @@ -5085,6 +5189,19 @@ export const makeClaudeAdapter = Effect.fn("makeClaudeAdapter")(function* (
const listSessions: ClaudeAdapterShape["listSessions"] = () =>
Effect.sync(() => Array.from(sessions.values(), ({ session }) => ({ ...session })));

// On-demand only. The SDK's token-count fallback can issue a model request,
// so this must never run automatically after turns (see #8610).
const getContextUsage: NonNullable<ClaudeAdapterShape["getContextUsage"]> = Effect.fn(
"getContextUsage",
)(function* (threadId) {
const context = yield* requireSession(threadId);
const response = yield* Effect.tryPromise({
try: () => context.query.getContextUsage(),
catch: (cause) => toRequestError(threadId, "context/usage", cause),
});
return toProviderContextUsage(response);
});

const hasSession: ClaudeAdapterShape["hasSession"] = (threadId) =>
Effect.sync(() => {
const context = sessions.get(threadId);
Expand Down Expand Up @@ -5135,6 +5252,7 @@ export const makeClaudeAdapter = Effect.fn("makeClaudeAdapter")(function* (
stopSession,
listSessions,
hasSession,
getContextUsage,
stopAll,
get streamEvents() {
return Stream.fromQueue(runtimeEventQueue);
Expand Down
Loading
Loading