Skip to content
Merged
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
17 changes: 17 additions & 0 deletions apps/desktop/src/lib/trpc/routers/ai-chat/index.ts
Original file line number Diff line number Diff line change
Expand Up @@ -88,6 +88,23 @@ export const createAiChatRouter = () => {
return { success: true };
}),

updateSessionConfig: publicProcedure
.input(
z.object({
sessionId: z.string(),
maxThinkingTokens: z.number().nullable().optional(),
model: z.string().nullable().optional(),
}),
)
.mutation(async ({ input }) => {
await chatSessionManager.updateAgentConfig({
sessionId: input.sessionId,
maxThinkingTokens: input.maxThinkingTokens,
model: input.model,
});
return { success: true };
}),

renameSession: publicProcedure
.input(
z.object({
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -309,6 +309,68 @@ export class ChatSessionManager extends EventEmitter {
await this.store.update(sessionId, patch);
}

async updateAgentConfig({
sessionId,
maxThinkingTokens,
model,
}: {
sessionId: string;
maxThinkingTokens?: number | null;
model?: string | null;
}): Promise<void> {
const session = this.sessions.get(sessionId);
if (!session) {
console.warn(
`[chat/session] Session ${sessionId} not found for config update`,
);
return;
}

const registration = this.provider.getAgentRegistration({
sessionId,
cwd: session.cwd,
});
Comment on lines +329 to +332
Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

⚠️ Potential issue | 🟠 Major

🧩 Analysis chain

🏁 Script executed:

#!/bin/bash
# First, let's look at the specific file and lines mentioned
cat -n apps/desktop/src/lib/trpc/routers/ai-chat/utils/session-manager/session-manager.ts | head -400 | tail -100

Repository: superset-sh/superset

Length of output: 2982


🏁 Script executed:

#!/bin/bash
# Find the updateAgentConfig method to understand context
rg -n 'updateAgentConfig' apps/desktop/src/lib/trpc/routers/ai-chat/utils/session-manager/session-manager.ts -B 10 -A 15

Repository: superset-sh/superset

Length of output: 733


🏁 Script executed:

#!/bin/bash
# Find startSession and restoreSession to see how they call getAgentRegistration
rg -n 'startSession|restoreSession' apps/desktop/src/lib/trpc/routers/ai-chat/utils/session-manager/session-manager.ts -B 5 -A 10

Repository: superset-sh/superset

Length of output: 617


🏁 Script executed:

#!/bin/bash
# Find the ActiveSession interface definition
rg -n 'interface ActiveSession' apps/desktop -A 20

Repository: superset-sh/superset

Length of output: 2236


🏁 Script executed:

#!/bin/bash
# Find the AgentProvider interface to understand getAgentRegistration signature
rg -n 'getAgentRegistration' apps/desktop -B 2 -A 10 --type=ts | head -100

Repository: superset-sh/superset

Length of output: 6883


updateAgentConfig rebuilds registration without paneId/tabId — routing context will be lost.

startSession and restoreSession pass paneId and tabId to getAgentRegistration, but updateAgentConfig omits them. Since ActiveSession does not store these values, they cannot be retrieved when the registration is rebuilt and POSTed back, causing any pane/tab routing information to be lost.

Store paneId and tabId in the ActiveSession interface so they're available when updating the agent config.

🤖 Prompt for AI Agents
In
`@apps/desktop/src/lib/trpc/routers/ai-chat/utils/session-manager/session-manager.ts`
around lines 327 - 330, updateAgentConfig rebuilds the agent registration
without paneId/tabId, causing routing context loss; add paneId and tabId fields
to the ActiveSession interface and ensure they are populated where sessions are
created (e.g., startSession and restoreSession) so updateAgentConfig can pass
session.paneId and session.tabId into getAgentRegistration when reconstructing
the registration; update references in updateAgentConfig, startSession,
restoreSession, and any session creation logic to set and read these new
ActiveSession properties.


if (maxThinkingTokens !== undefined) {
if (maxThinkingTokens === null) {
delete registration.bodyTemplate.maxThinkingTokens;
} else {
registration.bodyTemplate.maxThinkingTokens = maxThinkingTokens;
}
}

if (model !== undefined) {
if (model === null) {
delete registration.bodyTemplate.model;
} else {
registration.bodyTemplate.model = model;
}
}

const headers = buildProxyHeaders();
const res = await fetch(`${PROXY_URL}/v1/sessions/${sessionId}/agents`, {
method: "POST",
headers,
body: JSON.stringify({ agents: [registration] }),
});
if (!res.ok) {
throw new Error(
`POST /v1/sessions/${sessionId}/agents failed: ${res.status}`,
);
}

console.log(
`[chat/session] Updated agent config for ${sessionId}`,
[
maxThinkingTokens !== undefined &&
`maxThinkingTokens=${maxThinkingTokens}`,
model !== undefined && `model=${model}`,
]
.filter(Boolean)
.join(", "),
);
}
Comment on lines +312 to +372
Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

⚠️ Potential issue | 🟡 Minor

Missing error handling around fetch — inconsistent with other methods.

startSession and restoreSession wrap their fetch calls in try/catch and emit error events. updateAgentConfig lets fetch exceptions propagate as raw Errors, which will surface as INTERNAL_SERVER_ERROR in tRPC without contextual information. Consider wrapping in try/catch for consistency and better error messages.

As per coding guidelines, "Never swallow errors silently; at minimum log errors with context before rethrowing or handling them explicitly."

Proposed error handling
 	async updateAgentConfig({
 		sessionId,
 		maxThinkingTokens,
 	}: {
 		sessionId: string;
 		maxThinkingTokens?: number | null;
 	}): Promise<void> {
 		const session = this.sessions.get(sessionId);
 		if (!session) {
 			console.warn(
 				`[chat/session] Session ${sessionId} not found for config update`,
 			);
 			return;
 		}
 
 		const registration = this.provider.getAgentRegistration({
 			sessionId,
 			cwd: session.cwd,
 		});
 
 		if (maxThinkingTokens !== undefined) {
 			if (maxThinkingTokens === null) {
 				delete registration.bodyTemplate.maxThinkingTokens;
 			} else {
 				registration.bodyTemplate.maxThinkingTokens = maxThinkingTokens;
 			}
 		}
 
-		const headers = buildProxyHeaders();
-		const res = await fetch(`${PROXY_URL}/v1/sessions/${sessionId}/agents`, {
-			method: "POST",
-			headers,
-			body: JSON.stringify({ agents: [registration] }),
-		});
-		if (!res.ok) {
-			throw new Error(
-				`POST /v1/sessions/${sessionId}/agents failed: ${res.status}`,
-			);
-		}
-
-		console.log(
-			`[chat/session] Updated agent config for ${sessionId}`,
-			maxThinkingTokens !== undefined
-				? `maxThinkingTokens=${maxThinkingTokens}`
-				: "",
-		);
+		try {
+			const headers = buildProxyHeaders();
+			const res = await fetch(`${PROXY_URL}/v1/sessions/${sessionId}/agents`, {
+				method: "POST",
+				headers,
+				body: JSON.stringify({ agents: [registration] }),
+			});
+			if (!res.ok) {
+				throw new Error(
+					`POST /v1/sessions/${sessionId}/agents failed: ${res.status}`,
+				);
+			}
+
+			console.log(
+				`[chat/session] Updated agent config for ${sessionId}`,
+				maxThinkingTokens !== undefined
+					? `maxThinkingTokens=${maxThinkingTokens}`
+					: "",
+			);
+		} catch (error) {
+			const message = error instanceof Error ? error.message : String(error);
+			console.error(`[chat/session] Failed to update agent config:`, message);
+			throw error;
+		}
 	}
🤖 Prompt for AI Agents
In
`@apps/desktop/src/lib/trpc/routers/ai-chat/utils/session-manager/session-manager.ts`
around lines 312 - 358, The updateAgentConfig method currently performs fetch
without try/catch; wrap the network call and response check in a try/catch
inside updateAgentConfig to mirror startSession/restoreSession: use
buildProxyHeaders and fetch as before, on any thrown error or non-ok response
call this.emit('error', { sessionId, action: 'updateAgentConfig', error }) and
log a contextual message (including sessionId and maxThinkingTokens), then
rethrow a new Error with contextual details (or rethrow the caught error) so
callers get consistent, informative errors; keep existing registration logic
(this.provider.getAgentRegistration) and response handling (res.ok check) but
move them inside the try block.


isSessionActive(sessionId: string): boolean {
return this.sessions.has(sessionId);
}
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -16,6 +16,7 @@ import {
} from "@superset/ui/ai-elements/prompt-input";
import { Shimmer } from "@superset/ui/ai-elements/shimmer";
import { Suggestion, Suggestions } from "@superset/ui/ai-elements/suggestion";
import { ThinkingToggle } from "@superset/ui/ai-elements/thinking-toggle";
import { useCallback, useEffect, useRef, useState } from "react";
import {
HiMiniAtSymbol,
Expand Down Expand Up @@ -48,6 +49,9 @@ export function ChatInterface({
}: ChatInterfaceProps) {
const [selectedModel, setSelectedModel] = useState<ModelOption>(MODELS[1]);
const [modelSelectorOpen, setModelSelectorOpen] = useState(false);
const [thinkingEnabled, setThinkingEnabled] = useState(false);

const updateConfig = electronTrpc.aiChat.updateSessionConfig.useMutation();
Comment on lines +52 to +54
Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

⚠️ Potential issue | 🟡 Minor

thinkingEnabled is not initialized from the server-side session config.

If a session already has maxThinkingTokens set (e.g., after a restore), the toggle will show as disabled even though extended thinking is active on the backend. Consider initializing from the existing session or agent config to keep UI in sync.

🤖 Prompt for AI Agents
In
`@apps/desktop/src/renderer/screens/main/components/WorkspaceView/ContentView/TabsContent/TabView/ChatPane/ChatInterface/ChatInterface.tsx`
around lines 52 - 54, The thinkingEnabled local state is never initialized from
the server-side session/agent config so the toggle can be out of sync; update
ChatInterface to derive initial value from the session/agent config (check
session?.config?.maxThinkingTokens or agent?.config?.maxThinkingTokens) and call
setThinkingEnabled(Boolean(maxThinkingTokens)) in a useEffect that runs when the
session/agent changes, and ensure the toggle handler uses updateConfig.mutate to
set maxThinkingTokens to a default positive number when enabling and to 0 when
disabling so UI and backend stay synchronized.


const { data: config } = electronTrpc.aiChat.getConfig.useQuery();

Expand All @@ -60,6 +64,7 @@ export function ChatInterface({
stop,
addToolApprovalResponse,
connect,
collections,
} = useDurableChat({
sessionId,
proxyUrl: config?.proxyUrl ?? "http://localhost:8080",
Expand Down Expand Up @@ -214,6 +219,28 @@ export function ChatInterface({
[addToolApprovalResponse],
);

const handleThinkingToggle = useCallback(
(enabled: boolean) => {
setThinkingEnabled(enabled);
updateConfig.mutate({
sessionId,
maxThinkingTokens: enabled ? 10000 : null,
});
},
[sessionId, updateConfig],
);
Comment on lines +222 to +231
Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

⚠️ Potential issue | 🟡 Minor

Magic number 10000 and missing error handling on config mutation.

Two issues in this handler:

  1. Magic number: 10000 should be extracted to a named constant (e.g., DEFAULT_MAX_THINKING_TOKENS in constants.ts). As per coding guidelines, hardcoded magic numbers should be extracted to named constants at module top.

  2. No error handling: If updateConfig.mutate fails, the local thinkingEnabled state will be out of sync with the server. Add onError to roll back the optimistic state or at minimum log the failure — the coding guidelines require errors not be swallowed silently.

Proposed fix

In constants.ts (or top of file):

const DEFAULT_MAX_THINKING_TOKENS = 10_000;

Then in the handler:

 const handleThinkingToggle = useCallback(
   (enabled: boolean) => {
     setThinkingEnabled(enabled);
     updateConfig.mutate(
       {
         sessionId,
-        maxThinkingTokens: enabled ? 10000 : null,
+        maxThinkingTokens: enabled ? DEFAULT_MAX_THINKING_TOKENS : null,
+      },
+      {
+        onError: (err) => {
+          console.error("[chat/thinking] Failed to update config:", err);
+          setThinkingEnabled(!enabled);
+        },
       },
     );
   },
   [sessionId, updateConfig],
 );
🤖 Prompt for AI Agents
In
`@apps/desktop/src/renderer/screens/main/components/WorkspaceView/ContentView/TabsContent/TabView/ChatPane/ChatInterface/ChatInterface.tsx`
around lines 222 - 231, The handler handleThinkingToggle uses a hardcoded magic
number 10000 and calls updateConfig.mutate without error handling; extract the
magic number into a named constant (e.g., DEFAULT_MAX_THINKING_TOKENS in
constants.ts or at module top) and update the call to use that constant, and add
error handling for the mutation (use updateConfig.mutate or its equivalent with
an onError callback) to roll back or reset setThinkingEnabled when the mutation
fails and at minimum log the error so local state doesn't remain out of sync
with the server.


const handleModelSelect = useCallback(
(model: ModelOption) => {
setSelectedModel(model);
updateConfig.mutate({
sessionId,
model: model.id,
});
},
[sessionId, updateConfig],
);
Comment on lines +233 to +242
Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

⚠️ Potential issue | 🟡 Minor

handleModelSelect has no error rollback — optimistic state can drift.

Same pattern as handleThinkingToggle: if the mutation fails, selectedModel will reflect a model the backend didn't accept. Add onError to revert and log.

Proposed fix
 const handleModelSelect = useCallback(
   (model: ModelOption) => {
+    const previousModel = selectedModel;
     setSelectedModel(model);
-    updateConfig.mutate({
-      sessionId,
-      model: model.id,
-    });
+    updateConfig.mutate(
+      {
+        sessionId,
+        model: model.id,
+      },
+      {
+        onError: (err) => {
+          console.error("[chat/model] Failed to update model:", err);
+          setSelectedModel(previousModel);
+        },
+      },
+    );
   },
-  [sessionId, updateConfig],
+  [sessionId, updateConfig, selectedModel],
 );
📝 Committable suggestion

‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.

Suggested change
const handleModelSelect = useCallback(
(model: ModelOption) => {
setSelectedModel(model);
updateConfig.mutate({
sessionId,
model: model.id,
});
},
[sessionId, updateConfig],
);
const handleModelSelect = useCallback(
(model: ModelOption) => {
const previousModel = selectedModel;
setSelectedModel(model);
updateConfig.mutate(
{
sessionId,
model: model.id,
},
{
onError: (err) => {
console.error("[chat/model] Failed to update model:", err);
setSelectedModel(previousModel);
},
},
);
},
[sessionId, updateConfig, selectedModel],
);
🤖 Prompt for AI Agents
In
`@apps/desktop/src/renderer/screens/main/components/WorkspaceView/ContentView/TabsContent/TabView/ChatPane/ChatInterface/ChatInterface.tsx`
around lines 233 - 242, handleModelSelect currently performs an optimistic
update by calling setSelectedModel(model) before updateConfig.mutate and lacks
error rollback; change it to capture the previous model (e.g., const previous =
selectedModel) before setting the new model, pass an onError handler to
updateConfig.mutate that resets selectedModel back to previous and logs the
error, and optionally handle onSettled/onSuccess for consistency; ensure the
onError closure references the correct previous value and that the callback
still depends on sessionId and updateConfig.


const handleStop = useCallback(
(e: React.MouseEvent) => {
e.preventDefault();
Expand Down Expand Up @@ -296,15 +323,22 @@ export function ChatInterface({
<PromptInputButton>
<HiMiniAtSymbol className="size-4" />
</PromptInputButton>
<ThinkingToggle
enabled={thinkingEnabled}
onToggle={handleThinkingToggle}
/>
<ModelPicker
selectedModel={selectedModel}
onSelectModel={setSelectedModel}
onSelectModel={handleModelSelect}
open={modelSelectorOpen}
onOpenChange={setModelSelectorOpen}
/>
</PromptInputTools>
<div className="flex items-center gap-1">
<ContextIndicator />
<ContextIndicator
collections={collections}
modelId={selectedModel.id}
/>
<PromptInputSubmit
status={isLoading ? "streaming" : undefined}
onClick={isLoading ? handleStop : undefined}
Expand Down
Original file line number Diff line number Diff line change
@@ -1,36 +1,45 @@
import type { DurableChatCollections } from "@superset/durable-session/react";
import {
Context,
ContextCacheUsage,
ContextContent,
ContextContentBody,
ContextContentFooter,
ContextContentHeader,
ContextInputUsage,
ContextOutputUsage,
ContextReasoningUsage,
ContextTrigger,
} from "@superset/ui/ai-elements/context";
import { useLiveQuery } from "@tanstack/react-db";

const MOCK_CONTEXT = {
usedTokens: 84_200,
maxTokens: 200_000,
modelId: "claude-sonnet-4-5-20250929",
usage: {
inputTokens: 42_100,
outputTokens: 18_300,
totalTokens: 84_200,
reasoningTokens: 12_800,
cachedInputTokens: 11_000,
},
} as const;
const MAX_TOKENS = 200_000;

interface ContextIndicatorProps {
collections: DurableChatCollections;
modelId: string;
}

export function ContextIndicator({
collections,
modelId,
}: ContextIndicatorProps) {
const { data: statsRows } = useLiveQuery((q) =>
q.from({ s: collections.sessionStats }).select(({ s }) => ({ ...s })),
);

const stats = statsRows?.[0];
const usedTokens = stats?.totalTokens ?? 0;
const usage = {
inputTokens: stats?.promptTokens ?? 0,
outputTokens: stats?.completionTokens ?? 0,
totalTokens: stats?.totalTokens ?? 0,
};

export function ContextIndicator() {
return (
<Context
maxTokens={MOCK_CONTEXT.maxTokens}
modelId={MOCK_CONTEXT.modelId}
usage={MOCK_CONTEXT.usage}
usedTokens={MOCK_CONTEXT.usedTokens}
maxTokens={MAX_TOKENS}
modelId={modelId}
usage={usage}
usedTokens={usedTokens}
>
<ContextTrigger />
<ContextContent>
Expand All @@ -39,8 +48,6 @@ export function ContextIndicator() {
<div className="space-y-1">
<ContextInputUsage />
<ContextOutputUsage />
<ContextReasoningUsage />
<ContextCacheUsage />
</div>
</ContextContentBody>
<ContextContentFooter />
Expand Down
15 changes: 13 additions & 2 deletions apps/streams/src/claude-agent.ts
Original file line number Diff line number Diff line change
Expand Up @@ -32,6 +32,8 @@ const agentRequestSchema = z.object({
cwd: z.string().optional(),
env: z.record(z.string(), z.string()).optional(),
notification: notificationSchema.optional(),
maxThinkingTokens: z.number().optional(),
model: z.string().optional(),
});

interface SessionEntry {
Expand Down Expand Up @@ -169,7 +171,15 @@ app.post("/", async (c) => {
);
}

const { messages, sessionId, cwd, env: agentEnv, notification } = parsed.data;
const {
messages,
sessionId,
cwd,
env: agentEnv,
notification,
maxThinkingTokens,
model,
} = parsed.data;

const latestUserMessage = messages?.filter((m) => m.role === "user").pop();

Expand Down Expand Up @@ -197,14 +207,15 @@ app.post("/", async (c) => {
options: {
...(claudeSessionId && { resume: claudeSessionId }),
...(cwd && { cwd }),
model: process.env.CLAUDE_MODEL ?? DEFAULT_MODEL,
model: model ?? process.env.CLAUDE_MODEL ?? DEFAULT_MODEL,
maxTurns: MAX_AGENT_TURNS,
includePartialMessages: true,
permissionMode: "bypassPermissions" as const,
...(binaryPath && { pathToClaudeCodeExecutable: binaryPath }),
env: queryEnv,
abortController,
...(hooks && { hooks }),
...(maxThinkingTokens !== undefined && { maxThinkingTokens }),
},
});

Expand Down
5 changes: 1 addition & 4 deletions packages/ui/src/components/ai-elements/bash-tool.tsx
Original file line number Diff line number Diff line change
Expand Up @@ -82,10 +82,7 @@ export const BashTool = ({
if (state === "input-streaming") {
return (
<div
className={cn(
"flex items-start gap-1.5 rounded-md py-0.5",
className,
)}
className={cn("flex items-start gap-1.5 rounded-md py-0.5", className)}
>
<div className="min-w-0 flex flex-1 items-center gap-1.5">
<div className="flex min-w-0 items-center gap-1.5 text-xs text-muted-foreground">
Expand Down
55 changes: 55 additions & 0 deletions packages/ui/src/components/ai-elements/thinking-toggle.tsx
Original file line number Diff line number Diff line change
@@ -0,0 +1,55 @@
"use client";

import { BrainIcon } from "lucide-react";
import type { ComponentProps } from "react";
import { cn } from "../../lib/utils";
import { Button } from "../ui/button";
import {
Tooltip,
TooltipContent,
TooltipProvider,
TooltipTrigger,
} from "../ui/tooltip";

export type ThinkingToggleProps = Omit<
ComponentProps<typeof Button>,
"onClick" | "onToggle"
> & {
enabled: boolean;
onToggle: (enabled: boolean) => void;
};

export const ThinkingToggle = ({
enabled,
onToggle,
className,
...props
}: ThinkingToggleProps) => (
<TooltipProvider>
<Tooltip>
<TooltipTrigger asChild>
<Button
type="button"
variant="ghost"
size="icon-sm"
className={cn(
enabled && "bg-accent text-accent-foreground",
className,
)}
onClick={() => onToggle(!enabled)}
{...props}
>
<BrainIcon className="size-4" />
<span className="sr-only">
{enabled ? "Extended thinking enabled" : "Enable extended thinking"}
</span>
</Button>
</TooltipTrigger>
<TooltipContent>
<p>
{enabled ? "Extended thinking enabled" : "Enable extended thinking"}
</p>
</TooltipContent>
</Tooltip>
</TooltipProvider>
);
5 changes: 1 addition & 4 deletions packages/ui/src/components/ai-elements/tool-call.tsx
Original file line number Diff line number Diff line change
Expand Up @@ -29,10 +29,7 @@ export const ToolCall = ({

return (
<div
className={cn(
"flex items-start gap-1.5 rounded-md py-0.5",
className,
)}
className={cn("flex items-start gap-1.5 rounded-md py-0.5", className)}
>
<div className="min-w-0 flex flex-1 items-center gap-1.5">
<div className="flex min-w-0 items-center gap-1.5 text-xs text-muted-foreground">
Expand Down
7 changes: 1 addition & 6 deletions packages/ui/src/components/ai-elements/tool-interrupted.tsx
Original file line number Diff line number Diff line change
Expand Up @@ -13,12 +13,7 @@ export const ToolInterrupted = ({
subtitle,
className,
}: ToolInterruptedProps) => (
<div
className={cn(
"flex items-center gap-1.5 rounded-md py-0.5",
className,
)}
>
<div className={cn("flex items-center gap-1.5 rounded-md py-0.5", className)}>
<span className="text-xs text-muted-foreground">
{toolName} interrupted
</span>
Expand Down