Fimeg/letta-code
1
0
Fork 0
Code Issues Pull Requests Actions 1 Packages Projects Releases Wiki Activity

refactor(cli): unify TUI and headless stream processing (#784)

Browse Source 
Co-authored-by: Letta <noreply@letta.com>
This commit is contained in:
Charles Packer
2026-02-02 09:59:59 -08:00
committed by GitHub
parent 7297e334f0
commit b9eaaa1b5d
4 changed files with 463 additions and 477 deletions
Download Patch File Download Diff File Expand all files Collapse all files

184
src/cli/App.tsx
View File

@@ -159,6 +159,7 @@ import {
setToolCallsRunning,
toLines,
} from "./helpers/accumulator";
import { classifyApprovals } from "./helpers/approvalClassification";
import { backfillBuffers } from "./helpers/backfill";
import {
type AdvancedDiffSuccess,
@@ -3478,65 +3479,13 @@ export default function App({
}
// Check permissions for all approvals (including fancy UI tools)
const approvalResults = await Promise.all(
approvalsToProcess.map(async (approvalItem) => {
// Check if approval is incomplete (missing name)
// Note: toolArgs can be empty string for tools with no arguments (e.g., EnterPlanMode)
if (!approvalItem.toolName) {
return {
approval: approvalItem,
permission: {
decision: "deny" as const,
reason:
"Tool call incomplete - missing name or arguments",
},
context: null,
};
}
const parsedArgs = safeJsonParseOr<Record<string, unknown>>(
approvalItem.toolArgs,
{},
);
const permission = await checkToolPermission(
approvalItem.toolName,
parsedArgs,
);
const context = await analyzeToolApproval(
approvalItem.toolName,
parsedArgs,
);
return { approval: approvalItem, permission, context };
}),
);
// Categorize approvals by permission decision
// Fancy UI tools should always go through their dialog, even if auto-allowed
const needsUserInput: typeof approvalResults = [];
const autoDenied: typeof approvalResults = [];
const autoAllowed: typeof approvalResults = [];
for (const ac of approvalResults) {
const { approval, permission } = ac;
let decision = permission.decision;
// Some tools always need user input regardless of yolo mode
if (
alwaysRequiresUserInput(approval.toolName) &&
decision === "allow"
) {
decision = "ask";
}
if (decision === "ask") {
needsUserInput.push(ac);
} else if (decision === "deny") {
autoDenied.push(ac);
} else {
// decision === "allow"
autoAllowed.push(ac);
}
}
const { needsUserInput, autoAllowed, autoDenied } =
await classifyApprovals(approvalsToProcess, {
getContext: analyzeToolApproval,
alwaysRequiresUserInput,
missingNameReason:
"Tool call incomplete - missing name or arguments",
});
// Precompute diffs for file edit tools before execution (both auto-allowed and needs-user-input)
// This is needed for inline approval UI to show diffs, and for post-approval rendering
@@ -5054,58 +5003,12 @@ export default function App({
}
// There are pending approvals - check permissions (respects yolo mode)
const approvalResults = await Promise.all(
existingApprovals.map(async (approvalItem) => {
if (!approvalItem.toolName) {
return {
approval: approvalItem,
permission: {
decision: "deny" as const,
reason: "Tool call incomplete - missing name",
},
context: null,
};
}
const parsedArgs = safeJsonParseOr<Record<string, unknown>>(
approvalItem.toolArgs,
{},
);
const permission = await checkToolPermission(
approvalItem.toolName,
parsedArgs,
);
const context = await analyzeToolApproval(
approvalItem.toolName,
parsedArgs,
);
return { approval: approvalItem, permission, context };
}),
);
// Categorize by permission decision
const needsUserInput: typeof approvalResults = [];
const autoAllowed: typeof approvalResults = [];
const autoDenied: typeof approvalResults = [];
for (const ac of approvalResults) {
const { approval, permission } = ac;
let decision = permission.decision;
if (
alwaysRequiresUserInput(approval.toolName) &&
decision === "allow"
) {
decision = "ask";
}
if (decision === "ask") {
needsUserInput.push(ac);
} else if (decision === "deny") {
autoDenied.push(ac);
} else {
autoAllowed.push(ac);
}
}
const { needsUserInput, autoAllowed, autoDenied } =
await classifyApprovals(existingApprovals, {
getContext: analyzeToolApproval,
alwaysRequiresUserInput,
missingNameReason: "Tool call incomplete - missing name",
});
// If any approvals need user input, show dialog
if (needsUserInput.length > 0) {
@@ -7680,33 +7583,12 @@ ${SYSTEM_REMINDER_CLOSE}
if (existingApprovals && existingApprovals.length > 0) {
// There are pending approvals - check permissions first (respects yolo mode)
const approvalResults = await Promise.all(
existingApprovals.map(async (approvalItem) => {
if (!approvalItem.toolName) {
return {
approval: approvalItem,
permission: {
decision: "deny" as const,
reason: "Tool call incomplete - missing name",
},
context: null,
};
}
const parsedArgs = safeJsonParseOr<Record<string, unknown>>(
approvalItem.toolArgs,
{},
);
const permission = await checkToolPermission(
approvalItem.toolName,
parsedArgs,
);
const context = await analyzeToolApproval(
approvalItem.toolName,
parsedArgs,
);
return { approval: approvalItem, permission, context };
}),
);
const { needsUserInput, autoAllowed, autoDenied } =
await classifyApprovals(existingApprovals, {
getContext: analyzeToolApproval,
alwaysRequiresUserInput,
missingNameReason: "Tool call incomplete - missing name",
});
// Check if user cancelled during permission check
if (
@@ -7723,32 +7605,6 @@ ${SYSTEM_REMINDER_CLOSE}
return { submitted: false };
}
// Categorize by permission decision
const needsUserInput: typeof approvalResults = [];
const autoAllowed: typeof approvalResults = [];
const autoDenied: typeof approvalResults = [];
for (const ac of approvalResults) {
const { approval, permission } = ac;
let decision = permission.decision;
// Some tools always need user input regardless of yolo mode
if (
alwaysRequiresUserInput(approval.toolName) &&
decision === "allow"
) {
decision = "ask";
}
if (decision === "ask") {
needsUserInput.push(ac);
} else if (decision === "deny") {
autoDenied.push(ac);
} else {
autoAllowed.push(ac);
}
}
// If all approvals can be auto-handled (yolo mode), process them immediately
if (needsUserInput.length === 0) {
// Precompute diffs for file edit tools before execution (both auto-allowed and needs-user-input)

136
src/cli/helpers/approvalClassification.ts Normal file
View File

@@ -0,0 +1,136 @@
import type { ApprovalContext } from "../../permissions/analyzer";
import { checkToolPermission, getToolSchema } from "../../tools/manager";
import { safeJsonParseOr } from "./safeJsonParse";
import type { ApprovalRequest } from "./streamProcessor";
type ToolPermission = Awaited<ReturnType<typeof checkToolPermission>>;
export type ClassifiedApproval<TContext = ApprovalContext | null> = {
approval: ApprovalRequest;
permission: ToolPermission;
context: TContext | null;
parsedArgs: Record<string, unknown>;
missingRequiredArgs?: string[];
denyReason?: string;
};
export type ApprovalClassification<TContext = ApprovalContext | null> = {
needsUserInput: ClassifiedApproval<TContext>[];
autoAllowed: ClassifiedApproval<TContext>[];
autoDenied: ClassifiedApproval<TContext>[];
};
export type ClassifyApprovalsOptions<TContext = ApprovalContext | null> = {
getContext?: (
toolName: string,
parsedArgs: Record<string, unknown>,
) => Promise<TContext>;
alwaysRequiresUserInput?: (toolName: string) => boolean;
treatAskAsDeny?: boolean;
denyReasonForAsk?: string;
missingNameReason?: string;
requireArgsForAutoApprove?: boolean;
missingArgsReason?: (missing: string[]) => string;
};
export async function getMissingRequiredArgs(
toolName: string,
parsedArgs: Record<string, unknown>,
): Promise<string[]> {
const schema = getToolSchema(toolName);
const required =
(schema?.input_schema?.required as string[] | undefined) || [];
return required.filter(
(key) => !(key in parsedArgs) || parsedArgs[key] == null,
);
}
export async function classifyApprovals<TContext = ApprovalContext | null>(
approvals: ApprovalRequest[],
opts: ClassifyApprovalsOptions<TContext> = {},
): Promise<ApprovalClassification<TContext>> {
const needsUserInput: ClassifiedApproval<TContext>[] = [];
const autoAllowed: ClassifiedApproval<TContext>[] = [];
const autoDenied: ClassifiedApproval<TContext>[] = [];
const denyReasonForAsk =
opts.denyReasonForAsk ?? "Tool requires approval (headless mode)";
const missingNameReason =
opts.missingNameReason ?? "Tool call incomplete - missing name";
for (const approval of approvals) {
const toolName = approval.toolName;
if (!toolName) {
autoDenied.push({
approval,
permission: { decision: "deny", reason: missingNameReason },
context: null,
parsedArgs: {},
denyReason: missingNameReason,
});
continue;
}
const parsedArgs = safeJsonParseOr<Record<string, unknown>>(
approval.toolArgs || "{}",
{},
);
const permission = await checkToolPermission(toolName, parsedArgs);
const context = opts.getContext
? await opts.getContext(toolName, parsedArgs)
: null;
let decision = permission.decision;
if (opts.alwaysRequiresUserInput?.(toolName) && decision === "allow") {
decision = "ask";
}
if (decision === "ask" && opts.treatAskAsDeny) {
autoDenied.push({
approval,
permission,
context,
parsedArgs,
denyReason: denyReasonForAsk,
});
continue;
}
if (decision === "allow" && opts.requireArgsForAutoApprove) {
const missingRequiredArgs = await getMissingRequiredArgs(
toolName,
parsedArgs,
);
if (missingRequiredArgs.length > 0) {
const denyReason = opts.missingArgsReason
? opts.missingArgsReason(missingRequiredArgs)
: `Missing required parameter${missingRequiredArgs.length > 1 ? "s" : ""}: ${missingRequiredArgs.join(", ")}`;
autoDenied.push({
approval,
permission,
context,
parsedArgs,
missingRequiredArgs,
denyReason,
});
continue;
}
}
const entry: ClassifiedApproval<TContext> = {
approval,
permission,
context,
parsedArgs,
};
if (decision === "ask") {
needsUserInput.push(entry);
} else if (decision === "deny") {
autoDenied.push(entry);
} else {
autoAllowed.push(entry);
}
}
return { needsUserInput, autoAllowed, autoDenied };
}

63
src/cli/helpers/stream.ts
View File

@@ -13,6 +13,7 @@ import {
markIncompleteToolsAsCancelled,
onChunk,
} from "./accumulator";
import type { ErrorInfo } from "./streamProcessor";
import { StreamProcessor } from "./streamProcessor";
export type ApprovalRequest = {
@@ -21,6 +22,27 @@ export type ApprovalRequest = {
toolArgs: string;
};
export type DrainStreamHookContext = {
chunk: LettaStreamingResponse;
shouldOutput: boolean;
errorInfo?: ErrorInfo;
updatedApproval?: ApprovalRequest;
streamProcessor: StreamProcessor;
};
export type DrainStreamHookResult = {
shouldOutput?: boolean;
shouldAccumulate?: boolean;
stopReason?: StopReasonType;
};
export type DrainStreamHook = (
ctx: DrainStreamHookContext,
) =>
| DrainStreamHookResult
| undefined
| Promise<DrainStreamHookResult | undefined>;
type DrainResult = {
stopReason: StopReasonType;
lastRunId?: string | null;
@@ -37,6 +59,7 @@ export async function drainStream(
refresh: () => void,
abortSignal?: AbortSignal,
onFirstMessage?: () => void,
onChunkProcessed?: DrainStreamHook,
): Promise<DrainResult> {
const startTime = performance.now();
@@ -130,7 +153,8 @@ export async function drainStream(
logTiming(`TTFT: ${formatDuration(ttft)} (from POST to first content)`);
}
const { shouldOutput } = streamProcessor.processChunk(chunk);
const { shouldOutput, errorInfo, updatedApproval } =
streamProcessor.processChunk(chunk);
// Check abort signal before processing - don't add data after interrupt
if (abortSignal?.aborted) {
@@ -140,10 +164,40 @@ export async function drainStream(
break;
}
if (shouldOutput) {
let shouldOutputChunk = shouldOutput;
let shouldAccumulate = shouldOutput;
if (onChunkProcessed) {
const hookResult = await onChunkProcessed({
chunk,
shouldOutput: shouldOutputChunk,
errorInfo,
updatedApproval,
streamProcessor,
});
if (hookResult?.shouldOutput !== undefined) {
shouldOutputChunk = hookResult.shouldOutput;
}
if (hookResult?.shouldAccumulate !== undefined) {
shouldAccumulate = hookResult.shouldAccumulate;
} else {
shouldAccumulate = shouldOutputChunk;
}
if (hookResult?.stopReason) {
stopReason = hookResult.stopReason;
}
} else {
shouldAccumulate = shouldOutputChunk;
}
if (shouldAccumulate) {
onChunk(buffers, chunk);
queueMicrotask(refresh);
}
if (stopReason) {
break;
}
}
} catch (e) {
// Handle stream errors (e.g., JSON parse errors from SDK, network issues)
@@ -270,6 +324,7 @@ export async function drainStream(
* @param refresh - Callback to refresh UI
* @param abortSignal - Optional abort signal for cancellation
* @param onFirstMessage - Optional callback to invoke on first message chunk
* @param onChunkProcessed - Optional hook to observe/override per-chunk behavior
* @returns Result with stop_reason, approval info, and timing
*/
export async function drainStreamWithResume(
@@ -278,6 +333,7 @@ export async function drainStreamWithResume(
refresh: () => void,
abortSignal?: AbortSignal,
onFirstMessage?: () => void,
onChunkProcessed?: DrainStreamHook,
): Promise<DrainResult> {
const overallStartTime = performance.now();
@@ -288,6 +344,7 @@ export async function drainStreamWithResume(
refresh,
abortSignal,
onFirstMessage,
onChunkProcessed,
);
// If stream ended without proper stop_reason and we have resume info, try once to reconnect
@@ -333,6 +390,8 @@ export async function drainStreamWithResume(
buffers,
refresh,
abortSignal,
undefined,
onChunkProcessed,
);
// Use the resume result (should have proper stop_reason now)