Skip to content
Open
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
4 changes: 2 additions & 2 deletions src/main/services/ai/agents/projectAgent.ts
Original file line number Diff line number Diff line change
Expand Up @@ -137,12 +137,12 @@ If a database connection is failing or dbt cannot connect:

1. First diagnose; do not repair by editing connection files.
2. Use \`studio_connections_test\` when available to test the configured connection.
3. If needed, run \`studio_cli_run_dbt\` with \`dbt debug\` to verify the failure mode.
3. If needed, run \`studio_cli_run_dbt\` with \`command: "debug"\` to verify the failure mode. Do not pass \`"dbt debug"\` as the command; the tool adds the dbt executable.
4. Use \`getDbtLogs\` and read-only file inspection to explain the likely cause.
5. Ask the user to fix the connection in the DBT Studio Connections UI, secure keytar-backed credentials, or the intended connection-management workflow.
6. Wait for user confirmation that the connection has been corrected before continuing with dbt execution.

If the failure appears to be caused by invalid credentials, unreachable host, wrong port, missing network access, expired token, missing env vars, or bad connection metadata, do not rewrite \`profiles.yml\`. Report the issue clearly and direct the user to fix the connection configuration.
If the failure appears to be caused by invalid credentials, unreachable host, wrong port, missing network access, expired token, missing env vars, or bad connection metadata, run \`studio_cli_run_dbt\` with \`command: "debug"\` when it can add evidence, then do not rewrite \`profiles.yml\`. Report the exact missing env var or failed connection field clearly and direct the user to fix the connection configuration.

## Available Tools

Expand Down
41 changes: 41 additions & 0 deletions src/main/services/ai/tools/dbt.tools.ts
Original file line number Diff line number Diff line change
Expand Up @@ -9,6 +9,7 @@ import * as fs from 'fs';
import * as path from 'path';
import SettingsService from '../../settings.service';
import AgentService from '../../agent.service';
import SecureStorageService from '../../secureStorage.service';
import { TerminalConfirmGate } from './terminalConfirmGate';

// Security constraints
Expand Down Expand Up @@ -57,6 +58,42 @@ function isAllowedCommand(command: string, allowedCommands: string[]): boolean {
});
}

async function buildDbtProcessEnv(
projectPath: string,
): Promise<Record<string, string | undefined>> {
const env = { ...process.env };
const profilesPath = path.join(projectPath, 'profiles.yml');

if (!fs.existsSync(profilesPath)) {
return env;
}

const profilesContent = await fs.promises.readFile(profilesPath, 'utf8');
const envVarRegex = /env_var\(\s*["']([^"']+)["']/g;
const envVarNames = new Set<string>();

for (
let match = envVarRegex.exec(profilesContent);
match !== null;
match = envVarRegex.exec(profilesContent)
) {
envVarNames.add(match[1]);
}

await Promise.all(
Array.from(envVarNames).map(async (envVarName) => {
if (envVarName in env) return;

const storedValue = await SecureStorageService.getCredential(envVarName);
if (storedValue) {
env[envVarName] = storedValue;
}
}),
);

return env;
}

Comment thread
coderabbitai[bot] marked this conversation as resolved.
export async function executeDbtCommand(opts: {
command: string;
select?: string;
Expand Down Expand Up @@ -146,9 +183,12 @@ export async function executeDbtCommand(opts: {
mainWindow.webContents.send('cli:output', `> ${displayCmd}\n`);
}

const env = await buildDbtProcessEnv(projectPath);

return await new Promise((resolve) => {
const child = spawn(dbtExe, args, {
cwd: projectPath,
env,
shell: false,
});

Expand Down Expand Up @@ -435,6 +475,7 @@ export const runDbtCommand = tool({
// Execute with timeout
const output = execFileSync(dbtExe, args, {
cwd: projectPath,
env: await buildDbtProcessEnv(projectPath),
encoding: 'utf-8',
timeout: COMMAND_TIMEOUT,
maxBuffer: 1024 * 1024 * 10, // 10 MB buffer
Expand Down
50 changes: 29 additions & 21 deletions src/main/services/ai/tools/studio/cli.tools.ts
Original file line number Diff line number Diff line change
Expand Up @@ -10,15 +10,41 @@ const PHASE_19_ALLOWED_COMMANDS = [
'run',
'test',
'compile',
'deps',
'seed',
'debug',
'docs generate',
'source freshness',
];
const ALLOWED_COMMAND_DESCRIPTION =
'Allowed dbt command: run, test, compile, deps, seed, debug, docs generate, or source freshness';

const studioCliRunDbtInputSchema = z.object({
command: z
.string()
.refine((command) => PHASE_19_ALLOWED_COMMANDS.includes(command), {
message: ALLOWED_COMMAND_DESCRIPTION,
})
.describe(ALLOWED_COMMAND_DESCRIPTION),
select: z
.string()
.optional()
.describe(
'Optional dbt selector (e.g., "my_model", "tag:daily", "my_model+")',
),
extraArgs: z
.string()
.optional()
.describe('Optional additional dbt arguments (for advanced use cases)'),
});

type StudioCliRunDbtInput = z.infer<typeof studioCliRunDbtInputSchema>;

export function createStudioCliTools(options: {
projectPath?: string;
conversationId?: number;
mainWindow?: BrowserWindow;
}) {
}): Record<string, any> {
const { projectPath, conversationId, mainWindow } = options;
const cliEnabled = isToolEnabled(STUDIO_CLI_RUN_DBT_FLAG);

Expand All @@ -30,26 +56,8 @@ export function createStudioCliTools(options: {
studio_cli_run_dbt: tool({
description:
'Run an approved dbt CLI command for the active project. Always requires explicit user approval before execution.',
inputSchema: z.object({
command: z
.enum(['run', 'test', 'compile', 'docs generate', 'source freshness'])
.describe(
'Allowed dbt command: run, test, compile, docs generate, or source freshness',
),
select: z
.string()
.optional()
.describe(
'Optional dbt selector (e.g., "my_model", "tag:daily", "my_model+")',
),
extraArgs: z
.string()
.optional()
.describe(
'Optional additional dbt arguments (for advanced use cases)',
),
}),
execute: async ({ command, select, extraArgs }) => {
inputSchema: studioCliRunDbtInputSchema as any,
execute: async ({ command, select, extraArgs }: StudioCliRunDbtInput) => {
const startedAt = Date.now();

try {
Expand Down
6 changes: 6 additions & 0 deletions src/renderer/App.css
Original file line number Diff line number Diff line change
Expand Up @@ -81,3 +81,9 @@ a:hover {
background-size: 14px 14px;
cursor: pointer;
}

@keyframes spin {
to {
transform: rotate(360deg);
}
}
81 changes: 74 additions & 7 deletions src/renderer/components/chat/MessageRenderer.tsx
Original file line number Diff line number Diff line change
Expand Up @@ -64,6 +64,49 @@ interface MessageRendererProps {
orderedParts?: StreamContentPart[];
}

type RenderToolStatus = ToolCallState['status'];

function normalizePersistedToolStatus(
status: string | undefined,
output?: unknown,
error?: unknown,
isStreaming = false,
): RenderToolStatus {
const normalizedStatus = String(status ?? '').toLowerCase();

if (
['done', 'completed', 'success', 'succeeded'].includes(normalizedStatus)
) {
return 'done';
}

if (
['error', 'failed', 'failure', 'cancelled', 'canceled'].includes(
normalizedStatus,
)
) {
return 'error';
}

if (
isStreaming &&
['running', 'pending', 'queued'].includes(normalizedStatus)
) {
return 'running';
}

if (error) return 'error';
if (output !== undefined && output !== null) return 'done';

// Persisted history is never actively running. A historical running/pending
// tool without output most likely failed or was interrupted before result.
if (['running', 'pending', 'queued'].includes(normalizedStatus)) {
return 'error';
}

return 'done';
}

/**
* Converts persisted DB tool calls into AgentStep[] for rendering with AgentStepBlock.
* Tool calls are grouped by _stepNumber stored in toolInput.
Expand All @@ -90,18 +133,17 @@ function buildStepsFromToolCalls(
);
}

const statusMap: Record<string, 'done' | 'error'> = {
completed: 'done',
failed: 'error',
};

const toolCallState: ToolCallState = {
id: toolCallId,
toolName: tc.toolName,
args: displayArgs as Record<string, unknown>,
result: tc.toolOutput,
error: tc.errorMessage ?? undefined,
status: statusMap[tc.status] ?? 'done',
status: normalizePersistedToolStatus(
tc.status,
tc.toolOutput,
tc.errorMessage,
),
};

if (!stepMap.has(stepNumber)) {
Expand Down Expand Up @@ -594,6 +636,22 @@ export const MessageRenderer: React.FC<MessageRendererProps> = ({
return buildStepsFromToolCalls(toolCalls);
}, [toolCalls]);

const persistedToolStatusById = React.useMemo(() => {
const statusById = new Map<string, RenderToolStatus>();
if (!toolCalls) return statusById;

toolCalls.forEach((tc) => {
const toolCallId: string =
(tc.toolInput?.tcId as string) ?? String(tc.id);
statusById.set(
toolCallId,
normalizePersistedToolStatus(tc.status, tc.toolOutput, tc.errorMessage),
);
});

return statusById;
}, [toolCalls]);

// Fallback: if content looks like HTML (legacy TipTap), strip tags for display
const displayContent = React.useMemo(() => {
const looksLikeHtml = /<\/?[a-z][\s\S]*>/i.test(content || '');
Expand Down Expand Up @@ -687,6 +745,15 @@ export const MessageRenderer: React.FC<MessageRendererProps> = ({
}
// tool-call part
const tc = part as ToolCallContentPart;
const persistedStatus = persistedToolStatusById.get(tc.toolCallId);
const status =
persistedStatus ??
normalizePersistedToolStatus(
tc.status,
tc.result,
tc.error,
isStreaming,
);
return (
<ToolCallRow
key={tc.toolCallId}
Expand All @@ -696,7 +763,7 @@ export const MessageRenderer: React.FC<MessageRendererProps> = ({
args: tc.args,
result: tc.result,
error: tc.error,
status: tc.status,
status,
durationMs: tc.durationMs,
}}
/>
Expand Down
Loading
Loading