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
63 changes: 48 additions & 15 deletions src/app/admin/components/SessionTraceViewer.tsx
Original file line number Diff line number Diff line change
Expand Up @@ -8,18 +8,23 @@ import { Input } from '@/components/ui/input';
import { Button } from '@/components/ui/button';
import { Card, CardContent, CardHeader, CardTitle, CardDescription } from '@/components/ui/card';
import { Alert, AlertDescription } from '@/components/ui/alert';
import { MessageBubble } from '@/components/cloud-agent/MessageBubble';
import { MessageErrorBoundary } from '@/components/cloud-agent/MessageErrorBoundary';
import { MessageBubble as V1MessageBubble } from '@/components/cloud-agent/MessageBubble';
import { MessageErrorBoundary as V1MessageErrorBoundary } from '@/components/cloud-agent/MessageErrorBoundary';
import { convertToCloudMessages } from '@/components/cloud-agent/store/db-session-atoms';
import { MessageBubble as V2MessageBubble } from '@/components/cloud-agent-next/MessageBubble';
import { MessageErrorBoundary as V2MessageErrorBoundary } from '@/components/cloud-agent-next/MessageErrorBoundary';
import { isNewSession } from '@/lib/cloud-agent/session-type';
import {
useAdminSessionTrace,
useAdminSessionMessages,
useAdminApiConversationHistory,
} from '@/app/admin/api/session-traces/hooks';
import { Search, User, Calendar, Globe, GitBranch, Loader2, Download } from 'lucide-react';
import type { CloudMessage, Message } from '@/components/cloud-agent/types';
import type { StoredMessage } from '@/components/cloud-agent-next/types';

const UUID_REGEX = /^[0-9a-f]{8}-[0-9a-f]{4}-[0-9a-f]{4}-[0-9a-f]{4}-[0-9a-f]{12}$/i;
const SES_PREFIX = 'ses_';

function convertToMessage(cloudMessage: CloudMessage): Message & {
say?: string;
Expand Down Expand Up @@ -55,7 +60,10 @@ export function SessionTraceViewer() {

// Initialize from URL parameter on mount
useEffect(() => {
if (sessionIdFromUrl && UUID_REGEX.test(sessionIdFromUrl)) {
if (
sessionIdFromUrl &&
(UUID_REGEX.test(sessionIdFromUrl) || sessionIdFromUrl.startsWith(SES_PREFIX))
) {
setInputValue(sessionIdFromUrl);
setSearchedSessionId(sessionIdFromUrl);
}
Expand All @@ -71,9 +79,9 @@ export function SessionTraceViewer() {
setValidationError('Please enter a session ID');
return;
}
if (!UUID_REGEX.test(trimmed)) {
if (!UUID_REGEX.test(trimmed) && !trimmed.startsWith(SES_PREFIX)) {
setValidationError(
'Invalid UUID format. Expected format: xxxxxxxx-xxxx-xxxx-xxxx-xxxxxxxxxxxx'
'Invalid session ID. Expected a UUID (xxxxxxxx-xxxx-xxxx-xxxx-xxxxxxxxxxxx) or a v2 ID (ses_...)'
);
return;
}
Expand Down Expand Up @@ -117,14 +125,25 @@ export function SessionTraceViewer() {
}
};

const messages = useMemo(() => {
if (!messagesQuery.data?.messages) return [];
const isV2 = searchedSessionId ? isNewSession(searchedSessionId) : false;

const v1Messages = useMemo(() => {
if (!messagesQuery.data?.messages || messagesQuery.data.format === 'v2') return [];
const cloudMessages = convertToCloudMessages(
messagesQuery.data.messages as Array<Record<string, unknown>>
);
return cloudMessages.map(convertToMessage);
}, [messagesQuery.data]);

const v2Messages = useMemo(() => {
if (!messagesQuery.data?.messages || messagesQuery.data.format !== 'v2') return [];
// Server-side Zod validates minimal shape; full StoredMessage structure is
// guaranteed by the session-ingest worker that originally created the data.
return messagesQuery.data.messages as unknown as StoredMessage[];
Copy link
Contributor

Choose a reason for hiding this comment

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

[WARNING]: as unknown as StoredMessage[] bypasses type safety.

The project coding style strongly discourages as casts since they can't be statically checked. The server-side Zod schema (SessionSnapshotSchema) validates a minimal shape ({ info: { id }, parts: [{ id }] }), but StoredMessage expects info: OpenCodeMessage and parts: Part[] which are much richer types — a mismatch at runtime would silently produce incorrect data rather than a type error.

Consider defining a Zod schema on the client side that validates the full StoredMessage shape (or at least the fields MessageBubble actually accesses), and use .parse() / .safeParse() so mismatches surface as runtime errors rather than silent corruption.

}, [messagesQuery.data]);

const messageCount = isV2 ? v2Messages.length : v1Messages.length;

const breadcrumbs = (
<BreadcrumbItem>
<BreadcrumbPage>Session Traces</BreadcrumbPage>
Expand All @@ -138,13 +157,13 @@ export function SessionTraceViewer() {
<CardHeader>
<CardTitle>Session Trace Viewer</CardTitle>
<CardDescription>
Enter a CLI session ID (UUID) to view the full session trace
Enter a CLI session ID (UUID or ses_...) to view the full session trace
</CardDescription>
</CardHeader>
<CardContent>
<div className="flex gap-2">
<Input
placeholder="e.g., 550e8400-e29b-41d4-a716-446655440000"
placeholder="e.g., 550e8400-e29b-41d4-a716-446655440000 or ses_abc123..."
value={inputValue}
onChange={e => setInputValue(e.target.value)}
onKeyDown={e => e.key === 'Enter' && handleSearch()}
Expand Down Expand Up @@ -204,6 +223,12 @@ export function SessionTraceViewer() {
<span className="font-mono text-sm">{sessionQuery.data.git_url}</span>
</div>
)}
{sessionQuery.data.git_branch && (
<div className="flex items-center gap-2">
<GitBranch className="text-muted-foreground h-4 w-4" />
<span className="font-mono text-sm">{sessionQuery.data.git_branch}</span>
</div>
)}
<div className="flex items-center gap-2">
<Calendar className="text-muted-foreground h-4 w-4" />
<span className="text-sm">
Expand Down Expand Up @@ -236,7 +261,7 @@ export function SessionTraceViewer() {
<Card>
<CardHeader>
<div className="flex items-center justify-between">
<CardTitle>Messages ({messages.length})</CardTitle>
<CardTitle>Messages ({messageCount})</CardTitle>
{messagesQuery.data?.messages && (
<Button variant="outline" size="sm" onClick={handleDownloadMessages}>
<Download className="mr-2 h-4 w-4" />
Expand All @@ -251,14 +276,22 @@ export function SessionTraceViewer() {
<Loader2 className="h-4 w-4 animate-spin" />
<span>Loading messages...</span>
</div>
) : messages.length === 0 ? (
) : messageCount === 0 ? (
<p className="text-muted-foreground">No messages in this session</p>
) : isV2 ? (
<div className="space-y-2">
{v2Messages.map((msg, index) => (
<V2MessageErrorBoundary key={`${msg.info.id}-${index}`}>
<V2MessageBubble message={msg} isStreaming={false} />
</V2MessageErrorBoundary>
))}
</div>
) : (
<div className="space-y-2">
{messages.map((msg, index) => (
<MessageErrorBoundary key={`${msg.role}-${msg.timestamp}-${index}`}>
<MessageBubble message={msg} isStreaming={false} />
</MessageErrorBoundary>
{v1Messages.map((msg, index) => (
<V1MessageErrorBoundary key={`${msg.role}-${msg.timestamp}-${index}`}>
<V1MessageBubble message={msg} isStreaming={false} />
</V1MessageErrorBoundary>
))}
</div>
)}
Expand Down
97 changes: 91 additions & 6 deletions src/routers/admin-router.ts
Original file line number Diff line number Diff line change
Expand Up @@ -9,8 +9,11 @@ import {
user_auth_provider,
modelStats,
cliSessions,
cli_sessions_v2,
credit_transactions,
} from '@/db/schema';
import { isNewSession } from '@/lib/cloud-agent/session-type';
import { fetchSessionSnapshot, type SessionMessage } from '@/lib/session-ingest-client';
import { adminAppBuilderRouter } from '@/routers/admin-app-builder-router';
import { adminDeploymentsRouter } from '@/routers/admin-deployments-router';
import { adminKiloclawInstancesRouter } from '@/routers/admin-kiloclaw-instances-router';
Expand Down Expand Up @@ -69,6 +72,14 @@ const DeleteNoteSchema = z.object({
note_id: z.string(),
});

const UUID_REGEX = /^[0-9a-f]{8}-[0-9a-f]{4}-[0-9a-f]{4}-[0-9a-f]{4}-[0-9a-f]{12}$/i;
const sessionIdSchema = z
.string()
.min(1)
.refine(s => UUID_REGEX.test(s) || s.startsWith('ses_'), {
message: 'Must be a UUID or ses_-prefixed session ID',
});

const ResetAPIKeySchema = z.object({
userId: z.string(),
});
Expand Down Expand Up @@ -873,8 +884,42 @@ export const adminRouter = createTRPCRouter({

sessionTraces: createTRPCRouter({
get: adminProcedure
.input(z.object({ session_id: z.string().uuid() }))
.input(z.object({ session_id: sessionIdSchema }))
.query(async ({ input }) => {
if (isNewSession(input.session_id)) {
// V2 session — query cli_sessions_v2
const [session] = await db
.select()
.from(cli_sessions_v2)
.where(eq(cli_sessions_v2.session_id, input.session_id))
.limit(1);

if (!session) {
throw new TRPCError({
code: 'NOT_FOUND',
message: 'Session not found',
});
}

const user = await findUserById(session.kilo_user_id);

return {
...session,
// Fields that don't exist in v2 — null them out so the UI can handle both shapes
last_mode: null,
last_model: null,
user: user
? {
id: user.id,
email: user.google_user_email,
name: user.google_user_name,
image: user.google_user_image_url,
}
: null,
};
}

// V1 session — original logic
const [session] = await db
.select()
.from(cliSessions)
Expand All @@ -892,6 +937,8 @@ export const adminRouter = createTRPCRouter({

return {
...session,
// V1 doesn't have git_branch — null it out for a consistent shape
git_branch: null,
user: user
? {
id: user.id,
Expand All @@ -904,8 +951,40 @@ export const adminRouter = createTRPCRouter({
}),

getMessages: adminProcedure
.input(z.object({ session_id: z.string().uuid() }))
.input(z.object({ session_id: sessionIdSchema }))
.query(async ({ input }) => {
if (isNewSession(input.session_id)) {
// V2 session — fetch messages from the session-ingest worker.
// We need the owner's kilo_user_id to generate a service token.
const [session] = await db
.select({ kilo_user_id: cli_sessions_v2.kilo_user_id })
.from(cli_sessions_v2)
.where(eq(cli_sessions_v2.session_id, input.session_id))
.limit(1);

if (!session) {
throw new TRPCError({
code: 'NOT_FOUND',
message: 'Session not found',
});
}

try {
const snapshot = await fetchSessionSnapshot(input.session_id, session.kilo_user_id);
return {
messages: snapshot?.messages ?? ([] satisfies SessionMessage[]),
format: 'v2' as const,
};
} catch (error) {
console.error('[SessionTraces] Failed to fetch v2 session snapshot', {
sessionId: input.session_id,
error,
});
return { messages: [], format: 'v2' as const };
}
}

// V1 session — original logic
const [session] = await db
.select({
ui_messages_blob_url: cliSessions.ui_messages_blob_url,
Expand All @@ -922,20 +1001,26 @@ export const adminRouter = createTRPCRouter({
}

if (!session.ui_messages_blob_url) {
return { messages: [] };
return { messages: [], format: 'v1' as const };
}

try {
const messages = await getBlobContent(session.ui_messages_blob_url);
return { messages: (messages as unknown[]) ?? [] };
return { messages: (messages as unknown[]) ?? [], format: 'v1' as const };
} catch {
return { messages: [] };
return { messages: [], format: 'v1' as const };
}
}),

getApiConversationHistory: adminProcedure
.input(z.object({ session_id: z.string().uuid() }))
.input(z.object({ session_id: sessionIdSchema }))
.query(async ({ input }) => {
if (isNewSession(input.session_id)) {
// V2 sessions have no separate raw API conversation history
return { history: null };
}

// V1 session — original logic
const [session] = await db
.select({
api_conversation_history_blob_url: cliSessions.api_conversation_history_blob_url,
Expand Down