From 287ed18e13ca08530df4d22688f18d458e977ec5 Mon Sep 17 00:00:00 2001 From: ashrafchowdury Date: Tue, 21 Jul 2026 16:25:34 +0600 Subject: [PATCH 1/4] Enhance context management by integrating harness capabilities and context window for token budgeting --- .../AgentChatSlice/AgentConversation.tsx | 12 +++- .../AgentChatSlice/assets/contextBudget.ts | 57 +------------------ .../components/ContextBudgetIndicator.tsx | 12 ++-- .../hooks/useAgentModelKeyStatus.ts | 14 ++++- .../agenta-entities/src/workflow/index.ts | 1 + .../src/workflow/state/index.ts | 1 + .../src/workflow/state/inspectMeta.ts | 16 ++++++ 7 files changed, 50 insertions(+), 63 deletions(-) diff --git a/web/oss/src/components/AgentChatSlice/AgentConversation.tsx b/web/oss/src/components/AgentChatSlice/AgentConversation.tsx index f047c3c8c3..ba1117295b 100644 --- a/web/oss/src/components/AgentChatSlice/AgentConversation.tsx +++ b/web/oss/src/components/AgentChatSlice/AgentConversation.tsx @@ -17,6 +17,8 @@ import { } from "@agenta/entities/session" import {markTraceAsFresh} from "@agenta/entities/trace" import { + contextWindowForModel, + harnessCapabilitiesAtomFamily, invalidateAgentCommittedRevisionCache, workflowBuildKitOverlayReadyAtomFamily, workflowMolecule, @@ -724,6 +726,14 @@ const AgentConversation = ({ const modelKey = useAgentModelKeyStatus(entityId) const modelBlocked = modelKey.gateActive + // Context-window denominator for the token-budget indicator: the SDK model catalog's own + // `context_window`, delivered on the (global) harness-capabilities document — never hardcoded. + const harnessCapabilities = useAtomValue(harnessCapabilitiesAtomFamily("")) + const contextMaxTokens = useMemo( + () => contextWindowForModel(harnessCapabilities, modelKey.harness, modelKey.model), + [harnessCapabilities, modelKey.harness, modelKey.model], + ) + // ── Playground-native onboarding ────────────────────────────────────────── // This chat panel IS the onboarding surface while the agent is ephemeral: the empty state shows the // "what do you want to build?" hero and the composer renders Create-agent / Continue-in-IDE controls @@ -2255,7 +2265,7 @@ const AgentConversation = ({ {!onboardingActive ? ( ) : null} diff --git a/web/oss/src/components/AgentChatSlice/assets/contextBudget.ts b/web/oss/src/components/AgentChatSlice/assets/contextBudget.ts index d2df764c13..a168804d60 100644 --- a/web/oss/src/components/AgentChatSlice/assets/contextBudget.ts +++ b/web/oss/src/components/AgentChatSlice/assets/contextBudget.ts @@ -13,65 +13,14 @@ * double-counts the resent history (turn 2's prompt re-includes turn 1) so it climbs far faster * than real occupancy and never drops. Included only for the comparison. * - * The model's max context window is not exposed by the backend today (the registry stores only - * cost), so `MODEL_CONTEXT_WINDOWS` is the single hardcoded source for the "/ max (%)" denominator. + * The max context window (`maxTokens`) comes from the harness model catalog — the SDK's own + * `context_window` per model, delivered to the frontend on the harness-capabilities document and + * resolved via `contextWindowForModel` (@agenta/entities/workflow). Nothing is hardcoded here. */ import type {UIMessage} from "ai" import {getMessageUsage} from "./trace" -/** - * Model id (substring) → total context window in tokens. Keys are matched case-insensitively, - * after stripping any `provider/` or `provider:` prefix, by exact then longest-substring match, so - * dated snapshots (`claude-opus-4-20250101`) resolve to their base entry. Keep values conservative - * and extend as models are added; a follow-up can source these from litellm `get_model_info()`. - */ -export const MODEL_CONTEXT_WINDOWS: Record = { - // OpenAI - "gpt-4o-mini": 128_000, - "gpt-4o": 128_000, - "gpt-4.1-mini": 1_047_576, - "gpt-4.1-nano": 1_047_576, - "gpt-4.1": 1_047_576, - "gpt-4-turbo": 128_000, - "o1-mini": 128_000, - o1: 200_000, - "o3-mini": 200_000, - o3: 200_000, - "o4-mini": 200_000, - // Anthropic - "claude-3-5-sonnet": 200_000, - "claude-3-5-haiku": 200_000, - "claude-3-7-sonnet": 200_000, - "claude-opus-4": 200_000, - "claude-sonnet-4": 200_000, - "claude-haiku-4": 200_000, - // Google Gemini - "gemini-1.5-pro": 2_097_152, - "gemini-1.5-flash": 1_048_576, - "gemini-2.0-flash": 1_048_576, - "gemini-2.5-pro": 1_048_576, - "gemini-2.5-flash": 1_048_576, -} - -/** - * Resolve the max context window for a model id. Returns `null` when the model is unknown, in which - * case the indicator shows the token count without a "/ max (%)" fraction. - */ -export function resolveModelContextWindow(model: string | null | undefined): number | null { - if (!model) return null - const id = model.toLowerCase() - const bare = id.replace(/^[a-z0-9._-]+[/:]/, "") - if (MODEL_CONTEXT_WINDOWS[bare]) return MODEL_CONTEXT_WINDOWS[bare] - let best: {key: string; win: number} | null = null - for (const [key, win] of Object.entries(MODEL_CONTEXT_WINDOWS)) { - if (bare.includes(key) && (!best || key.length > best.key.length)) { - best = {key, win} - } - } - return best?.win ?? null -} - export interface ContextBudget { /** Latest assistant turn's total tokens — current window occupancy. `null` until a turn has usage. */ occupancyTokens: number | null diff --git a/web/oss/src/components/AgentChatSlice/components/ContextBudgetIndicator.tsx b/web/oss/src/components/AgentChatSlice/components/ContextBudgetIndicator.tsx index 2b2dcbb996..ca1c5cbb91 100644 --- a/web/oss/src/components/AgentChatSlice/components/ContextBudgetIndicator.tsx +++ b/web/oss/src/components/AgentChatSlice/components/ContextBudgetIndicator.tsx @@ -14,13 +14,14 @@ import {useMemo} from "react" import type {UIMessage} from "ai" import {Tooltip, Typography} from "antd" -import {computeContextBudget, resolveModelContextWindow} from "../assets/contextBudget" +import {computeContextBudget} from "../assets/contextBudget" const {Text} = Typography export interface ContextBudgetIndicatorProps { messages: readonly UIMessage[] - model: string | null + /** Max context window for the selected model, from the harness catalog; `null` when unknown. */ + maxTokens: number | null className?: string } @@ -33,11 +34,8 @@ const fmt = (n: number): string => { const pctLabel = (pct: number | null): string => (pct == null ? "" : ` (${Math.round(pct * 100)}%)`) -const ContextBudgetIndicator = ({messages, model, className}: ContextBudgetIndicatorProps) => { - const budget = useMemo(() => { - const maxTokens = resolveModelContextWindow(model) - return computeContextBudget(messages, maxTokens) - }, [messages, model]) +const ContextBudgetIndicator = ({messages, maxTokens, className}: ContextBudgetIndicatorProps) => { + const budget = useMemo(() => computeContextBudget(messages, maxTokens), [messages, maxTokens]) if (budget.occupancyTokens == null && budget.runningSumTokens == null) return null diff --git a/web/oss/src/components/AgentChatSlice/hooks/useAgentModelKeyStatus.ts b/web/oss/src/components/AgentChatSlice/hooks/useAgentModelKeyStatus.ts index d34d944318..e9fd8f2e87 100644 --- a/web/oss/src/components/AgentChatSlice/hooks/useAgentModelKeyStatus.ts +++ b/web/oss/src/components/AgentChatSlice/hooks/useAgentModelKeyStatus.ts @@ -15,6 +15,8 @@ export interface AgentModelKeyStatus { provider: string | null /** The selected model id (display). */ model: string | null + /** The selected harness type (e.g. "pi_core" / "claude"), from `agent.harness.kind`. */ + harness: string | null /** Whether the project's vault holds a key for that provider. */ hasKey: boolean /** The canonical vault provider entry for the model's provider (to open the configure drawer). */ @@ -40,6 +42,10 @@ interface LlmRef { connection?: {mode?: unknown} | null } +interface HarnessRef { + kind?: unknown +} + /** * Model → provider → vault-key detection for an agent. The `agent.llm` value is a structured ModelRef * carrying its `provider`; we check the project's vault (`standardSecretsAtom`) for a key for that @@ -64,8 +70,13 @@ export function useAgentModelKeyStatus(entityId: string): AgentModelKeyStatus { const keySetupDone = useAtomValue(providerKeySetupDoneAtom) return useMemo(() => { - const llm = (config as {agent?: {llm?: LlmRef}} | null)?.agent?.llm + const agent = (config as {agent?: {llm?: LlmRef; harness?: HarnessRef}} | null)?.agent + const llm = agent?.llm const model = typeof llm?.model === "string" && llm.model ? llm.model : null + const harness = + typeof agent?.harness?.kind === "string" && agent.harness.kind + ? agent.harness.kind + : null // Provider is stored on the ModelRef; fall back to a `provider/id` model prefix (Pi naming). const provider = typeof llm?.provider === "string" && llm.provider @@ -90,6 +101,7 @@ export function useAgentModelKeyStatus(entityId: string): AgentModelKeyStatus { return { provider, model, + harness, hasKey: !!providerEntry?.key, providerEntry, loading, diff --git a/web/packages/agenta-entities/src/workflow/index.ts b/web/packages/agenta-entities/src/workflow/index.ts index 409b879943..663ea84fdc 100644 --- a/web/packages/agenta-entities/src/workflow/index.ts +++ b/web/packages/agenta-entities/src/workflow/index.ts @@ -54,6 +54,7 @@ export { // Per-harness capability map from the `/inspect` response `meta` (agent playground picker). export { harnessCapabilitiesAtomFamily, + contextWindowForModel, type HarnessCapabilities, type HarnessCapabilitiesMap, type ModelCatalogEntry, diff --git a/web/packages/agenta-entities/src/workflow/state/index.ts b/web/packages/agenta-entities/src/workflow/state/index.ts index 9e80c04bc7..57e3959bd5 100644 --- a/web/packages/agenta-entities/src/workflow/state/index.ts +++ b/web/packages/agenta-entities/src/workflow/state/index.ts @@ -16,6 +16,7 @@ export {workflowMolecule, type WorkflowMolecule} from "./molecule" export { harnessCapabilitiesAtomFamily, + contextWindowForModel, type HarnessCapabilities, type HarnessCapabilitiesMap, type ModelCatalogEntry, diff --git a/web/packages/agenta-entities/src/workflow/state/inspectMeta.ts b/web/packages/agenta-entities/src/workflow/state/inspectMeta.ts index 19e0a13a7c..81d25a8aef 100644 --- a/web/packages/agenta-entities/src/workflow/state/inspectMeta.ts +++ b/web/packages/agenta-entities/src/workflow/state/inspectMeta.ts @@ -126,3 +126,19 @@ export const harnessCapabilitiesAtomFamily = atomFamily((_harnessRef: string) => return query.data ?? null }), ) + +/** + * The model's context window (max input tokens) from the harness catalog, matched by exact id — the + * same `id`-keyed join the model picker uses. `null` when the catalog, harness, or entry is absent, + * or the entry carries no `context_window`. Source of truth is the SDK model catalog + * (`model_catalog.py`), so no window is ever hardcoded on the frontend. + */ +export function contextWindowForModel( + capabilities: HarnessCapabilitiesMap | null | undefined, + harness: string | null | undefined, + modelId: string | null | undefined, +): number | null { + if (!capabilities || !harness || !modelId) return null + const entry = capabilities[harness]?.model_catalog?.find((e) => e.id === modelId) + return entry?.context_window ?? null +} From 6306223a6ce4d033a58a3d08653dd4dd57c907a7 Mon Sep 17 00:00:00 2001 From: ashrafchowdury Date: Tue, 21 Jul 2026 16:26:47 +0600 Subject: [PATCH 2/4] Refactor ContextBudget and ContextBudgetIndicator for improved clarity and performance --- .../AgentChatSlice/assets/contextBudget.ts | 39 ++----- .../components/ContextBudgetIndicator.tsx | 102 +++++++++++------- 2 files changed, 74 insertions(+), 67 deletions(-) diff --git a/web/oss/src/components/AgentChatSlice/assets/contextBudget.ts b/web/oss/src/components/AgentChatSlice/assets/contextBudget.ts index a168804d60..efe11ff019 100644 --- a/web/oss/src/components/AgentChatSlice/assets/contextBudget.ts +++ b/web/oss/src/components/AgentChatSlice/assets/contextBudget.ts @@ -1,17 +1,11 @@ /** - * Context token budget — the data behind the "Context 31.3k / 200k (16%)" indicator. + * Context token budget — the data behind the ambient "Context 39% used" meter above the composer. * * This slice runs each turn through `useChat`; the runner stamps real token usage onto every * assistant message (`metadata.usage = {input, output, total, cost}` → read via `getMessageUsage`). - * From those per-turn totals we compute TWO candidate measures so the UI can show them side by - * side until we pick one: - * - * - **occupancy** = the LATEST assistant turn's total tokens. Because the whole conversation is - * resent every turn, this already grows with the chat AND equals how full the context window is - * right now — so it predicts compaction and correctly DROPS after a compaction/summarization. - * - **runningSum** = Σ of every assistant turn's total tokens. A cumulative usage meter; it - * double-counts the resent history (turn 2's prompt re-includes turn 1) so it climbs far faster - * than real occupancy and never drops. Included only for the comparison. + * Occupancy = the LATEST assistant turn's total tokens: because the whole conversation is resent + * every turn, this equals how full the context window is right now, so it predicts compaction and + * correctly DROPS after a compaction/summarization. * * The max context window (`maxTokens`) comes from the harness model catalog — the SDK's own * `context_window` per model, delivered to the frontend on the harness-capabilities document and @@ -24,23 +18,14 @@ import {getMessageUsage} from "./trace" export interface ContextBudget { /** Latest assistant turn's total tokens — current window occupancy. `null` until a turn has usage. */ occupancyTokens: number | null - /** Σ of every assistant turn's total tokens this session. `null` until a turn has usage. */ - runningSumTokens: number | null - /** Per-turn totals, oldest → newest (for the tooltip / debugging). */ - perTurnTokens: number[] /** Model context window, or `null` when unknown. */ maxTokens: number | null /** occupancy / max, clamped 0..1; `null` when max unknown or no usage yet. */ occupancyPct: number | null - /** runningSum / max, clamped 0..1; `null` when max unknown or no usage yet. */ - runningSumPct: number | null } -const pctOf = (n: number | null, max: number | null): number | null => - n != null && max ? Math.min(n / max, 1) : null - /** - * Compute both budget measures from a session's messages and its model's context window. + * Compute window occupancy from a session's messages and its model's context window. * Reads real per-turn usage off assistant messages (`getMessageUsage`), preferring `totalTokens` * and falling back to `promptTokens` when only the input side was stamped. */ @@ -48,24 +33,18 @@ export function computeContextBudget( messages: readonly UIMessage[], maxTokens: number | null, ): ContextBudget { - const perTurnTokens: number[] = [] + let occupancyTokens: number | null = null for (const message of messages) { if (message.role !== "assistant") continue const usage = getMessageUsage(message) const total = usage?.totalTokens ?? usage?.promptTokens - if (typeof total === "number") perTurnTokens.push(total) + if (typeof total === "number") occupancyTokens = total } - const hasUsage = perTurnTokens.length > 0 - const occupancyTokens = hasUsage ? perTurnTokens[perTurnTokens.length - 1] : null - const runningSumTokens = hasUsage ? perTurnTokens.reduce((a, b) => a + b, 0) : null - return { occupancyTokens, - runningSumTokens, - perTurnTokens, maxTokens, - occupancyPct: pctOf(occupancyTokens, maxTokens), - runningSumPct: pctOf(runningSumTokens, maxTokens), + occupancyPct: + occupancyTokens != null && maxTokens ? Math.min(occupancyTokens / maxTokens, 1) : null, } } diff --git a/web/oss/src/components/AgentChatSlice/components/ContextBudgetIndicator.tsx b/web/oss/src/components/AgentChatSlice/components/ContextBudgetIndicator.tsx index ca1c5cbb91..b7253d74d4 100644 --- a/web/oss/src/components/AgentChatSlice/components/ContextBudgetIndicator.tsx +++ b/web/oss/src/components/AgentChatSlice/components/ContextBudgetIndicator.tsx @@ -1,23 +1,21 @@ /** - * ContextBudgetIndicator — the compact "Context 31.3k / 200k (16%)" strip above the composer. + * ContextBudgetIndicator — the ambient "Context 39% used" meter above the composer. * - * v1 shows BOTH candidate measures side by side so we can compare them on a live agent and pick - * the right one (then drop the other + its label): - * - **Ctx** = occupancy (latest turn's tokens = how full the window is now; predicts compaction) - * - **Σ** = running sum (cumulative tokens across the session) + * A slim fill bar + one plain-language line. Quiet (neutral) while there's room, escalating to + * amber then red with alarm wording as the window fills, so non-technical users get a signal + * without reading token math. Exact token counts live in the tooltip for power users. * - * Renders nothing until at least one turn has reported usage. When the model's context window is - * unknown it shows the token counts without the "/ max (%)" fraction. + * Renders nothing until a turn has reported usage. When the model's context window is unknown it + * falls back to a quiet raw token count with no bar or percent. */ import {useMemo} from "react" +import {Warning} from "@phosphor-icons/react" import type {UIMessage} from "ai" -import {Tooltip, Typography} from "antd" +import {Tooltip} from "antd" import {computeContextBudget} from "../assets/contextBudget" -const {Text} = Typography - export interface ContextBudgetIndicatorProps { messages: readonly UIMessage[] /** Max context window for the selected model, from the harness catalog; `null` when unknown. */ @@ -25,58 +23,88 @@ export interface ContextBudgetIndicatorProps { className?: string } -/** Compact token formatting: 31_300 → "31.3k", 1_047_576 → "1.0M". */ +/** Compact token formatting, no noisy trailing ".0": 34_200 → "34.2k", 407_027 → "407k", 1_048_576 → "1M". */ +const compact = (n: number, div: number, suffix: string): string => { + const s = (n / div).toFixed(1) + return `${s.endsWith(".0") ? s.slice(0, -2) : s}${suffix}` +} const fmt = (n: number): string => { - if (n >= 1_000_000) return `${(n / 1_000_000).toFixed(1)}M` - if (n >= 1_000) return `${(n / 1_000).toFixed(1)}k` + if (n >= 1_000_000) return compact(n, 1_000_000, "M") + if (n >= 1_000) return compact(n, 1_000, "k") return `${n}` } -const pctLabel = (pct: number | null): string => (pct == null ? "" : ` (${Math.round(pct * 100)}%)`) +type Tier = "normal" | "warn" | "danger" + +const tierFor = (pct: number | null): Tier => + pct == null ? "normal" : pct >= 0.9 ? "danger" : pct >= 0.75 ? "warn" : "normal" + +const FILL: Record = { + normal: "bg-colorTextTertiary", + warn: "bg-colorWarning", + danger: "bg-colorError", +} + +const TEXT: Record = { + normal: "text-colorTextTertiary", + warn: "text-colorWarning", + danger: "text-colorError", +} const ContextBudgetIndicator = ({messages, maxTokens, className}: ContextBudgetIndicatorProps) => { const budget = useMemo(() => computeContextBudget(messages, maxTokens), [messages, maxTokens]) - if (budget.occupancyTokens == null && budget.runningSumTokens == null) return null + const {occupancyTokens: occ, maxTokens: max, occupancyPct: pct} = budget + if (occ == null) return null - const maxSuffix = budget.maxTokens ? ` / ${fmt(budget.maxTokens)}` : "" + const tier = tierFor(pct) + const pctInt = pct != null ? Math.round(pct * 100) : null + const label = + pctInt != null + ? tier === "danger" + ? `Context almost full · ${pctInt}%` + : `Context ${pctInt}% used` + : `Context ${fmt(occ)} tokens` return ( - Ctx — tokens in the context window right now (latest turn). Predicts - when the conversation compacts, and drops after it does. - - - Σ — total tokens across the whole session (cumulative usage). + How full the model's context window is this turn. When it fills up, + older messages get summarized to make room. - {budget.maxTokens - ? `Model window ≈ ${fmt(budget.maxTokens)} tokens` - : "Model context window unknown"} + {max + ? `${fmt(occ)} / ${fmt(max)} tokens` + : `${fmt(occ)} tokens · model window unknown`} } placement="topRight" mouseEnterDelay={0.4} > - - {budget.occupancyTokens != null ? ( - - Ctx {fmt(budget.occupancyTokens)} - {maxSuffix} - {pctLabel(budget.occupancyPct)} - - ) : null} - {budget.runningSumTokens != null ? ( - - · Σ {fmt(budget.runningSumTokens)} - {pctLabel(budget.runningSumPct)} + + {pctInt != null ? ( + + ) : null} - + + {tier === "danger" ? ( + + ) : null} + {label} + + ) } From d6dca3ced7a8380d04feca8fef67bf7d8caba923 Mon Sep 17 00:00:00 2001 From: ashrafchowdury Date: Tue, 21 Jul 2026 16:40:37 +0600 Subject: [PATCH 3/4] Update tooltip text in ContextBudgetIndicator for clarity on context window behavior --- .../AgentChatSlice/components/ContextBudgetIndicator.tsx | 6 ++++-- 1 file changed, 4 insertions(+), 2 deletions(-) diff --git a/web/oss/src/components/AgentChatSlice/components/ContextBudgetIndicator.tsx b/web/oss/src/components/AgentChatSlice/components/ContextBudgetIndicator.tsx index b7253d74d4..d508f61d4d 100644 --- a/web/oss/src/components/AgentChatSlice/components/ContextBudgetIndicator.tsx +++ b/web/oss/src/components/AgentChatSlice/components/ContextBudgetIndicator.tsx @@ -71,8 +71,10 @@ const ContextBudgetIndicator = ({messages, maxTokens, className}: ContextBudgetI title={
- How full the model's context window is this turn. When it fills up, - older messages get summarized to make room. + Context window: +
+ As it fills up, earlier messages are compressed into summaries, which may + reduce recall of older details. {max From 40ca548c1ae86654c0cef1c96c3b541d5d46e7a0 Mon Sep 17 00:00:00 2001 From: ashrafchowdury Date: Tue, 21 Jul 2026 16:54:18 +0600 Subject: [PATCH 4/4] Enhance danger tier styling in ContextBudgetIndicator for improved visibility --- .../AgentChatSlice/components/ContextBudgetIndicator.tsx | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/web/oss/src/components/AgentChatSlice/components/ContextBudgetIndicator.tsx b/web/oss/src/components/AgentChatSlice/components/ContextBudgetIndicator.tsx index d508f61d4d..16b687d1e2 100644 --- a/web/oss/src/components/AgentChatSlice/components/ContextBudgetIndicator.tsx +++ b/web/oss/src/components/AgentChatSlice/components/ContextBudgetIndicator.tsx @@ -42,13 +42,13 @@ const tierFor = (pct: number | null): Tier => const FILL: Record = { normal: "bg-colorTextTertiary", warn: "bg-colorWarning", - danger: "bg-colorError", + danger: "bg-colorError opacity-70", } const TEXT: Record = { normal: "text-colorTextTertiary", warn: "text-colorWarning", - danger: "text-colorError", + danger: "text-colorError opacity-70", } const ContextBudgetIndicator = ({messages, maxTokens, className}: ContextBudgetIndicatorProps) => {