feat assistant context and agent trace

2026-08-14 11:32:11 +02:00
1 parent aea2832370
commit f4f6773f26
@@ -0,0 +1,12 @@
+import { groupPartByType } from "@assistant-ui/react";
+
+/**
+ * Keep model reasoning and ordinary tool calls in one trace. Human tools and
+ * other standalone tool UIs stay outside the trace so approvals and questions
+ * retain the full conversation width.
+ */
+export const assistantMessageGroupBy = groupPartByType({
+  reasoning: ["group-agent-trace", "group-reasoning"],
+  "tool-call": ["group-agent-trace", "group-tool"],
+  "standalone-tool-call": [],
+});
@@ -1,6 +1,7 @@
 import { describe, expect, test } from "vitest";

 import {
+  getAssistantContextBreakdown,
  getAssistantConversationUsage,
  getAssistantTelemetry,
  getAssistantUsage,
@@ -262,3 +263,53 @@ describe("getAssistantConversationUsage", () => {
    });
  });
 });
+
+describe("getAssistantContextBreakdown", () => {
+  test("combines the latest context reading with the latest reported token categories", () => {
+    const usage = getAssistantConversationUsage([
+      createUsageMetadata({
+        requestId: "reported",
+        inputTokens: 12_777,
+        outputTokens: 55,
+        cost: 0.00041,
+        contextUsed: 13_343,
+      }),
+      {
+        custom: {
+          telemetry: {
+            requestId: "streaming",
+            provider: "openrouter",
+            modelId: "example/model",
+            startedAt: "2026-08-04T10:00:02.000Z",
+            steps: [],
+          },
+        },
+      },
+    ]);
+
+    expect(getAssistantContextBreakdown(usage)).toMatchObject({
+      contextUsed: 13_343,
+      contextLength: 128_000,
+      remaining: 114_657,
+      percentage: 10,
+      inputTokens: 12_777,
+      outputTokens: 55,
+    });
+  });
+
+  test("does not report an invalid percentage for an unavailable capacity", () => {
+    expect(
+      getAssistantContextBreakdown({
+        turns: [],
+        inputTokens: 0,
+        outputTokens: 0,
+        totalTokens: 0,
+        cachedInputTokens: 0,
+        reasoningTokens: 0,
+        cost: 0,
+        contextUsed: 10,
+        contextLength: 0,
+      }),
+    ).toEqual({});
+  });
+});
@@ -122,6 +122,17 @@ export type AssistantConversationUsage = {
  contextLength?: number;
 };

+export type AssistantContextBreakdown = {
+  contextUsed?: number;
+  contextLength?: number;
+  remaining?: number;
+  percentage?: number;
+  inputTokens?: number;
+  outputTokens?: number;
+  cachedInputTokens?: number;
+  reasoningTokens?: number;
+};
+
 const getFiniteNonNegativeNumber = (value: unknown) =>
  typeof value === "number" && Number.isFinite(value) && value >= 0 ? value : undefined;

@@ -412,6 +423,35 @@ export const getAssistantConversationUsage = (messageMetadata: unknown[]): Assis
  };
 };

+export const getAssistantContextBreakdown = (usage: AssistantConversationUsage): AssistantContextBreakdown => {
+  const latestTokenTurn = usage.turns.findLast(
+    (turn) =>
+      turn.inputTokens !== undefined ||
+      turn.outputTokens !== undefined ||
+      turn.cachedInputTokens !== undefined ||
+      turn.reasoningTokens !== undefined,
+  );
+  const { contextUsed, contextLength } = usage;
+  const hasContext = contextUsed !== undefined && contextLength !== undefined && contextLength > 0;
+
+  return {
+    ...(hasContext
+      ? {
+          contextUsed,
+          contextLength,
+          remaining: Math.max(0, contextLength - contextUsed),
+          percentage: Math.min(100, Math.round((contextUsed / contextLength) * 100)),
+        }
+      : {}),
+    ...(latestTokenTurn?.inputTokens !== undefined ? { inputTokens: latestTokenTurn.inputTokens } : {}),
+    ...(latestTokenTurn?.outputTokens !== undefined ? { outputTokens: latestTokenTurn.outputTokens } : {}),
+    ...(latestTokenTurn?.cachedInputTokens !== undefined
+      ? { cachedInputTokens: latestTokenTurn.cachedInputTokens }
+      : {}),
+    ...(latestTokenTurn?.reasoningTokens !== undefined ? { reasoningTokens: latestTokenTurn.reasoningTokens } : {}),
+  };
+};
+
 export const resolveAssistantContextWindowTelemetry = (
  current: AssistantRequestTelemetry,
  previousMetadata: unknown[],
@@ -663,6 +663,17 @@
  text-align: start;
 }

+.tool[data-compact] {
+  margin-block: 0;
+  padding: 0.4rem 0.5rem;
+  border-radius: var(--mantine-radius-sm);
+  background: var(--mantine-color-body);
+}
+
+.tool[data-compact] .toolHeader > :global(.mantine-Group-root) {
+  min-width: 0;
+}
+
 .approvalPanel {
  margin-top: 0.8rem;
  padding-top: 0.8rem;
@@ -961,6 +972,10 @@
  min-width: 0;
 }

+.agentTraceTools {
+  min-width: 0;
+}
+
 .source {
  display: inline-flex;
  align-items: center;
@@ -1191,6 +1206,43 @@
  background: var(--mantine-color-default);
 }

+.contextBreakdown {
+  display: grid;
+  grid-template-columns: 4.75rem minmax(0, 1fr);
+  align-items: center;
+  gap: 0.6rem 0.75rem;
+  padding: 0.55rem;
+  border: 1px solid var(--mantine-color-default-border);
+  border-radius: var(--mantine-radius-sm);
+  background: var(--mantine-color-default);
+}
+
+.contextBreakdownRing {
+  flex: 0 0 auto;
+}
+
+.contextBreakdownWindow {
+  display: grid;
+  min-width: 0;
+  grid-template-columns: repeat(3, minmax(0, 1fr));
+  gap: 0.45rem;
+}
+
+.contextBreakdownWindow > div,
+.contextBreakdownCategories > div {
+  min-width: 0;
+}
+
+.contextBreakdownCategories {
+  display: grid;
+  min-width: 0;
+  grid-column: 1 / -1;
+  grid-template-columns: repeat(4, minmax(0, 1fr));
+  gap: 0.45rem;
+  padding-top: 0.5rem;
+  border-top: 1px solid var(--mantine-color-default-border);
+}
+
 .conversationTurn {
  min-width: 0;
  padding: 0.5rem 0.55rem;
@@ -1626,6 +1678,14 @@
  }

  .conversationUsageSummary {
+    grid-template-columns: repeat(3, minmax(0, 1fr));
+  }
+
+  .contextBreakdownWindow {
+    grid-template-columns: repeat(2, minmax(0, 1fr));
+  }
+
+  .contextBreakdownCategories {
    grid-template-columns: repeat(2, minmax(0, 1fr));
  }

@@ -27,6 +27,7 @@ import {
  ActionBarPrimitive,
  AttachmentPrimitive,
  BranchPickerPrimitive,
+  ChainOfThoughtByIndicesProvider,
  ChainOfThoughtPrimitive,
  ComposerPrimitive,
  ErrorPrimitive,
@@ -128,14 +129,20 @@ import { remarkAssistantDirectives, resolveAssistantDirectiveEntity } from "./as
 import type { AssistantDirectiveEntity } from "./assistant-markdown-directives";
 import { normalizeAssistantMarkdown } from "./assistant-markdown";
 import { getSafeAssistantAttachmentImageSource, getSafeAssistantMarkdownImageSource } from "./assistant-markdown-image";
-import { getAssistantConversationUsage, getAssistantTelemetry } from "./assistant-message-metadata";
-import type { AssistantConversationTurnUsage } from "./assistant-message-metadata";
+import {
+  getAssistantContextBreakdown,
+  getAssistantConversationUsage,
+  getAssistantTelemetry,
+} from "./assistant-message-metadata";
+import type { AssistantContextBreakdown, AssistantConversationTurnUsage } from "./assistant-message-metadata";
+import { assistantMessageGroupBy } from "./assistant-message-grouping";
 import type { AssistantPendingAction } from "./assistant-pending-action";
 import { AssistantQuestionPortalProvider } from "./assistant-question-portal";
 import type { AssistantReasoningMode, AssistantRuntimeModelOption } from "./assistant-preferences";
 import { useAssistantReasoningState } from "./assistant-reasoning-state";
 import { getNearestTriggerScrollTop } from "./assistant-trigger-scroll";
 import { getToolResultPresentation } from "./assistant-tool-result";
+import { getAssistantToolTraceTarget } from "./assistant-tool-trace";
 import { getSafeAssistantHttpUrl } from "./assistant-url";

 // Shiki's web grammar bundle is intentionally loaded only when a settled response contains a code
@@ -800,6 +807,8 @@ const ToolResultPreview = ({ result, toolName }: { result: unknown; toolName: st
  );
 };

+const AgentTraceToolContext = createContext(false);
+
 const ToolPart = ({
  toolCallId,
  toolName,
@@ -812,6 +821,7 @@ const ToolPart = ({
  timing,
 }: ToolCallMessagePartProps) => {
  const t = useScopedI18n("common.assistant");
+  const compact = useContext(AgentTraceToolContext);
  const [opened, setOpened] = useState(false);
  const [approvalResponse, setApprovalResponse] = useState<"approve" | "deny" | null>(null);
  const completed = status?.type === "complete";
@@ -823,6 +833,8 @@ const ToolPart = ({
      status?.type === "incomplete" ||
      (typeof result === "object" && result !== null && "error" in result));
  const successful = completed && !denied && !failed;
+  const compactPresentation = compact && !awaitingApproval;
+  const traceTarget = getAssistantToolTraceTarget(args);
  const duration = timing?.completedAt !== undefined ? Math.max(0, timing.completedAt - timing.startedAt) : undefined;
  const autoApprovalInProgress = useAssistantAutomaticAction({
    toolCallId,
@@ -835,7 +847,7 @@ const ToolPart = ({
  });

  return (
-    <Box className={classes.tool}>
+    <Box className={classes.tool} data-compact={compactPresentation || undefined}>
      <UnstyledButton
        className={classes.toolHeader}
        onClick={() => setOpened((current) => !current)}
@@ -851,9 +863,16 @@ const ToolPart = ({
            >
              {denied || failed ? <IconX size={13} /> : successful ? <IconCheck size={13} /> : <IconRobot size={13} />}
            </ThemeIcon>
-            <Text size="sm" fw={600}>
-              {toolName.replaceAll("_", " ")}
-            </Text>
+            <Box miw={0}>
+              <Text size={compactPresentation ? "xs" : "sm"} fw={650} lineClamp={1}>
+                {toolName.replaceAll("_", " ")}
+              </Text>
+              {compactPresentation && traceTarget && (
+                <Text size="xs" c="dimmed" lineClamp={1}>
+                  {traceTarget}
+                </Text>
+              )}
+            </Box>
          </Group>
          <Group gap={6} wrap="nowrap">
            {duration !== undefined && (
@@ -926,9 +945,14 @@ const ToolPart = ({
          )}
        </Box>
      )}
-      {successful && result !== undefined && <ToolResultPreview result={result} toolName={toolName} />}
+      {!compactPresentation && successful && result !== undefined && (
+        <ToolResultPreview result={result} toolName={toolName} />
+      )}
      <Collapse expanded={opened}>
        <Stack gap="xs" mt="sm">
+          {compactPresentation && successful && result !== undefined && (
+            <ToolResultPreview result={result} toolName={toolName} />
+          )}
          <Box>
            <Text size="xs" fw={600} c="dimmed">
              {t("toolInput")}
@@ -953,13 +977,21 @@ const ToolPart = ({
  );
 };

+const AgentTraceToolGroup = ({ children }: { children?: ReactNode }) => (
+  <AgentTraceToolContext.Provider value>
+    <Stack className={classes.agentTraceTools} gap={4}>
+      {children}
+    </Stack>
+  </AgentTraceToolContext.Provider>
+);
+
 const ChainOfThoughtLayout = ({ children }: { children?: ReactNode }) => (
-  <Stack className={classes.reasoningParts} gap="xs">
+  <Stack className={classes.reasoningParts} gap={6}>
    {children}
  </Stack>
 );

-const AssistantChainOfThought = () => {
+const AssistantChainOfThought = ({ children }: { children?: ReactNode }) => {
  const t = useScopedI18n("common.assistant");
  const contentId = useId();
  const { collapsed: preferredCollapsed, setCollapsed: setPreferredCollapsed } = useContext(ReasoningVisibilityContext);
@@ -988,9 +1020,7 @@ const AssistantChainOfThought = () => {
        </UnstyledButton>
      </ChainOfThoughtPrimitive.AccordionTrigger>
      <Box id={contentId} className={classes.reasoningContent} aria-busy={running} hidden={collapsed}>
-        <ChainOfThoughtPrimitive.Parts
-          components={{ Reasoning: ReasoningPart, tools: { Fallback: ToolPart }, Layout: ChainOfThoughtLayout }}
-        />
+        <ChainOfThoughtLayout>{children}</ChainOfThoughtLayout>
      </Box>
    </ChainOfThoughtPrimitive.Root>
  );
@@ -1512,6 +1542,64 @@ const ConversationTurn = ({
  );
 };

+const ConversationContextBreakdown = ({ breakdown }: { breakdown: AssistantContextBreakdown }) => {
+  const t = useScopedI18n("common.assistant");
+  const hasContext = breakdown.percentage !== undefined;
+  const categories = [
+    { label: t("usage.input"), value: breakdown.inputTokens },
+    { label: t("usage.cached"), value: breakdown.cachedInputTokens },
+    { label: t("usage.output"), value: breakdown.outputTokens },
+    { label: t("usage.reasoning"), value: breakdown.reasoningTokens },
+  ];
+
+  return (
+    <Box className={classes.contextBreakdown}>
+      <RingProgress
+        className={classes.contextBreakdownRing}
+        size={76}
+        thickness={7}
+        roundCaps
+        sections={
+          hasContext ? [{ value: breakdown.percentage ?? 0, color: getContextColor(breakdown.percentage ?? 0) }] : []
+        }
+        label={
+          <Text ta="center" size="xs" fw={700}>
+            {hasContext ? `${breakdown.percentage}%` : "–"}
+          </Text>
+        }
+      />
+      <Box className={classes.contextBreakdownWindow}>
+        {[
+          { label: t("usage.used"), value: breakdown.contextUsed },
+          { label: t("usage.remaining"), value: breakdown.remaining },
+          { label: t("usage.capacity"), value: breakdown.contextLength },
+        ].map((item) => (
+          <Box key={item.label}>
+            <Text size="xs" c="dimmed">
+              {item.label}
+            </Text>
+            <Text size="sm" fw={650}>
+              {item.value?.toLocaleString() ?? t("usage.notReported")}
+            </Text>
+          </Box>
+        ))}
+      </Box>
+      <Box className={classes.contextBreakdownCategories}>
+        {categories.map((category) => (
+          <Box key={category.label}>
+            <Text size="xs" c="dimmed">
+              {category.label}
+            </Text>
+            <Text size="xs" fw={650}>
+              {category.value?.toLocaleString() ?? t("usage.notReported")}
+            </Text>
+          </Box>
+        ))}
+      </Box>
+    </Box>
+  );
+};
+
 const ConversationContext = () => {
  const t = useScopedI18n("common.assistant");
  const [opened, setOpened] = useState(false);
@@ -1529,12 +1617,11 @@ const ConversationContext = () => {
    [messages],
  );
  const usage = useMemo(() => getAssistantConversationUsage(metadata), [metadata]);
+  const breakdown = useMemo(() => getAssistantContextBreakdown(usage), [usage]);
  const hasTokenUsage = usage.turns.some((turn) => turn.totalTokens !== undefined);
  const hasCost = usage.turns.some((turn) => turn.cost !== undefined);
-  const hasContext = usage.contextUsed !== undefined && usage.contextLength !== undefined;
-  const contextPercentage = hasContext
-    ? Math.min(100, Math.round(((usage.contextUsed ?? 0) / Math.max(usage.contextLength ?? 1, 1)) * 100))
-    : 0;
+  const hasContext = breakdown.percentage !== undefined;
+  const contextPercentage = breakdown.percentage ?? 0;
  const quickLabel = [
    hasTokenUsage ? `${usage.totalTokens.toLocaleString()} ${t("usage.tokens")}` : undefined,
    hasCost ? formatCost(usage.cost) : undefined,
@@ -1607,20 +1694,16 @@ const ConversationContext = () => {
            </div>
            <div>
              <Text size="xs" c="dimmed">
-                {t("usage.used")}
+                {t("usage.agentSteps")}
              </Text>
              <Text size="sm" fw={650}>
-                {usage.contextUsed?.toLocaleString() ?? t("usage.notReported")}
+                {usage.turns.reduce((total, turn) => total + turn.telemetry.steps.length, 0).toLocaleString()}
              </Text>
-              {usage.contextLength !== undefined && (
-                <Text size="xs" c="dimmed">
-                  / {usage.contextLength.toLocaleString()}
-                </Text>
-              )}
            </div>
          </Box>
+          <ConversationContextBreakdown breakdown={breakdown} />
          <Divider />
-          <ScrollArea.Autosize mah="min(21rem, 48dvh)" type="auto" offsetScrollbars>
+          <ScrollArea.Autosize mah="min(14rem, 30dvh)" type="auto" offsetScrollbars>
            <Stack gap={5}>
              {usage.turns.length === 0 ? (
                <Text size="sm" c="dimmed" py="xs">
@@ -1739,17 +1822,44 @@ const WebSearchActivity = () => {

 const AssistantMessage = () => (
  <MessagePrimitive.Root className={`${classes.message} ${classes.assistantMessage}`}>
-    <MessagePrimitive.Parts
-      unstable_showEmptyOnNonTextEnd={false}
-      components={{
-        Empty: AssistantMessagePending,
-        Text: AssistantTextPart,
-        Source: SourcePart,
-        File: FilePart,
-        Image: ImagePart,
-        ChainOfThought: AssistantChainOfThought,
+    <MessagePrimitive.GroupedParts groupBy={assistantMessageGroupBy} indicator="empty">
+      {({ part, children }) => {
+        switch (part.type) {
+          case "group-agent-trace": {
+            const startIndex = part.indices[0];
+            const endIndex = part.indices.at(-1);
+            if (startIndex === undefined || endIndex === undefined) return null;
+            return (
+              <ChainOfThoughtByIndicesProvider startIndex={startIndex} endIndex={endIndex}>
+                <AssistantChainOfThought>{children}</AssistantChainOfThought>
+              </ChainOfThoughtByIndicesProvider>
+            );
+          }
+          case "group-reasoning":
+            return <>{children}</>;
+          case "group-tool":
+            return <AgentTraceToolGroup>{children}</AgentTraceToolGroup>;
+          case "indicator":
+            return <AssistantMessagePending status={{ type: "running" }} />;
+          case "text":
+            return <AssistantTextPart />;
+          case "reasoning":
+            return <ReasoningPart {...part} />;
+          case "tool-call":
+            return part.toolUI ?? <ToolPart {...part} />;
+          case "source":
+            return <SourcePart {...part} />;
+          case "file":
+            return <FilePart {...part} />;
+          case "image":
+            return <ImagePart {...part} />;
+          case "data":
+            return part.dataRendererUI;
+          default:
+            return null;
+        }
      }}
-    />
+    </MessagePrimitive.GroupedParts>
    <RuntimeError />
    <WebSearchActivity />
    <AssistantMessageActions />
@@ -1,10 +1,12 @@
 // @vitest-environment jsdom

+import type { ReactNode } from "react";
 import { act, createElement, useEffect } from "react";
 import { createRoot } from "react-dom/client";
 import type { AssistantClient, ChatModelAdapter, ThreadMessageLike } from "@assistant-ui/react";
 import {
  AssistantRuntimeProvider,
+  ChainOfThoughtByIndicesProvider,
  ChainOfThoughtPrimitive,
  MessagePrimitive,
  ThreadPrimitive,
@@ -13,6 +15,7 @@ import {
 } from "@assistant-ui/react";
 import { afterEach, describe, expect, test } from "vitest";

+import { assistantMessageGroupBy } from "./assistant-message-grouping";
 import { useAssistantReasoningState } from "./assistant-reasoning-state";

 (globalThis as typeof globalThis & { IS_REACT_ACT_ENVIRONMENT: boolean }).IS_REACT_ACT_ENVIRONMENT = true;
@@ -27,7 +30,13 @@ afterEach(() => {
 const ReasoningPart = ({ text }: { text: string }) => <span data-testid="reasoning-part">{text}</span>;
 const ToolPart = ({ toolName }: { toolName: string }) => <span data-testid="tool-part">{toolName}</span>;

-const TestChainOfThought = ({ preferredCollapsed = true }: { preferredCollapsed?: boolean }) => {
+const TestChainOfThought = ({
+  preferredCollapsed = true,
+  children,
+}: {
+  preferredCollapsed?: boolean;
+  children?: ReactNode;
+}) => {
  const { chainStatus, collapsed } = useAssistantReasoningState(preferredCollapsed);

  return (
@@ -36,14 +45,38 @@ const TestChainOfThought = ({ preferredCollapsed = true }: { preferredCollapsed?
      data-status={chainStatus.type}
      data-collapsed={String(collapsed)}
    >
-      <ChainOfThoughtPrimitive.Parts components={{ Reasoning: ReasoningPart, tools: { Fallback: ToolPart } }} />
+      {children}
    </ChainOfThoughtPrimitive.Root>
  );
 };

 const TestMessage = () => (
  <MessagePrimitive.Root>
-    <MessagePrimitive.Parts components={{ ChainOfThought: TestChainOfThought }} />
+    <MessagePrimitive.GroupedParts groupBy={assistantMessageGroupBy} indicator="never">
+      {({ part, children }) => {
+        switch (part.type) {
+          case "group-agent-trace": {
+            const startIndex = part.indices[0];
+            const endIndex = part.indices.at(-1);
+            if (startIndex === undefined || endIndex === undefined) return null;
+            return (
+              <ChainOfThoughtByIndicesProvider startIndex={startIndex} endIndex={endIndex}>
+                <TestChainOfThought>{children}</TestChainOfThought>
+              </ChainOfThoughtByIndicesProvider>
+            );
+          }
+          case "group-reasoning":
+          case "group-tool":
+            return <>{children}</>;
+          case "reasoning":
+            return <ReasoningPart text={part.text} />;
+          case "tool-call":
+            return <ToolPart toolName={part.toolName} />;
+          default:
+            return null;
+        }
+      }}
+    </MessagePrimitive.GroupedParts>
  </MessagePrimitive.Root>
 );

@@ -0,0 +1,20 @@
+import { describe, expect, test } from "vitest";
+
+import { getAssistantToolTraceTarget } from "./assistant-tool-trace";
+
+describe("getAssistantToolTraceTarget", () => {
+  test("uses a human-readable tool argument instead of an opaque id", () => {
+    expect(getAssistantToolTraceTarget({ id: "app-1", name: "  Wikipedia  " })).toBe("Wikipedia");
+    expect(getAssistantToolTraceTarget({ integrationIds: ["integration-1"], searchText: " discord   icon " })).toBe(
+      "discord icon",
+    );
+  });
+
+  test("does not expose arbitrary argument values in the collapsed trace", () => {
+    expect(getAssistantToolTraceTarget({ id: "secret-looking-id", apiKey: "do-not-show" })).toBeNull();
+  });
+
+  test("bounds long labels so the trace remains compact", () => {
+    expect(getAssistantToolTraceTarget({ title: "a".repeat(100) })).toBe(`${"a".repeat(77)}…`);
+  });
+});
@@ -0,0 +1,16 @@
+const targetKeys = ["name", "title", "query", "searchText", "appName", "widgetName"] as const;
+
+export const getAssistantToolTraceTarget = (args: unknown): string | null => {
+  if (!args || typeof args !== "object" || Array.isArray(args)) return null;
+  const record = args as Record<string, unknown>;
+
+  for (const key of targetKeys) {
+    const value = record[key];
+    if (typeof value !== "string") continue;
+    const normalized = value.trim().replaceAll(/\s+/gu, " ");
+    if (normalized.length === 0) continue;
+    return normalized.length > 80 ? `${normalized.slice(0, 77)}…` : normalized;
+  }
+
+  return null;
+};