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
2 changes: 1 addition & 1 deletion src/entrypoints/cli.tsx
Original file line number Diff line number Diff line change
Expand Up @@ -146,7 +146,7 @@ async function main(): Promise<void> {
shutdown1PEventLogging,
logForDebugging,
registerPermissionHandler(server, handler) {
server.setNotificationHandler(ChannelPermissionRequestNotificationSchema() as any, async notification =>
server.setNotificationHandler(ChannelPermissionRequestNotificationSchema(), async notification =>
handler(notification.params),
);
},
Expand Down
5 changes: 3 additions & 2 deletions src/hooks/useIdeLogging.ts
Original file line number Diff line number Diff line change
Expand Up @@ -3,9 +3,10 @@ import { logEvent } from 'src/services/analytics/index.js'
import { z } from 'zod/v4'
import type { MCPServerConnection } from '../services/mcp/types.js'
import { getConnectedIdeClient } from '../utils/ide.js'
import type { AnyObjectSchema } from '@modelcontextprotocol/sdk/server/zod-compat.js'
import { lazySchema } from '../utils/lazySchema.js'

const LogEventSchema = lazySchema(() =>
const LogEventSchema: () => AnyObjectSchema = lazySchema(() =>
z.object({
method: z.literal('log_event'),
params: z.object({
Expand All @@ -27,7 +28,7 @@ export function useIdeLogging(mcpClients: MCPServerConnection[]): void {
if (ideClient) {
// Register the log event handler
ideClient.client.setNotificationHandler(
LogEventSchema() as any,
LogEventSchema(),
notification => {
const { eventName, eventData } = notification.params
logEvent(
Expand Down
5 changes: 3 additions & 2 deletions src/hooks/useIdeSelection.ts
Original file line number Diff line number Diff line change
Expand Up @@ -6,6 +6,7 @@ import type {
MCPServerConnection,
} from '../services/mcp/types.js'
import { getConnectedIdeClient } from '../utils/ide.js'
import type { AnyObjectSchema } from '@modelcontextprotocol/sdk/server/zod-compat.js'
import { lazySchema } from '../utils/lazySchema.js'
export type SelectionPoint = {
line: number
Expand All @@ -29,7 +30,7 @@ export type IDESelection = {
}

// Define the selection changed notification schema
const SelectionChangedSchema = lazySchema(() =>
const SelectionChangedSchema: () => AnyObjectSchema = lazySchema(() =>
z.object({
method: z.literal('selection_changed'),
params: z.object({
Expand Down Expand Up @@ -110,7 +111,7 @@ export function useIdeSelection(

// Register notification handler for selection_changed events
ideClient.client.setNotificationHandler(
SelectionChangedSchema() as any,
SelectionChangedSchema(),
notification => {
if (currentIDERef.current !== ideClient) {
return
Expand Down
5 changes: 3 additions & 2 deletions src/hooks/usePromptsFromClaudeInChrome.tsx
Original file line number Diff line number Diff line change
Expand Up @@ -6,11 +6,12 @@ import { callIdeRpc } from '../services/mcp/client.js';
import type { ConnectedMCPServer, MCPServerConnection } from '../services/mcp/types.js';
import type { PermissionMode } from '../types/permissions.js';
import { CLAUDE_IN_CHROME_MCP_SERVER_NAME, isTrackedClaudeInChromeTabId } from '../utils/claudeInChrome/common.js';
import type { AnyObjectSchema } from '@modelcontextprotocol/sdk/server/zod-compat.js';
import { lazySchema } from '../utils/lazySchema.js';
import { enqueuePendingNotification } from '../utils/messageQueueManager.js';

// Schema for the prompt notification from Chrome extension (JSON-RPC 2.0 format)
const ClaudeInChromePromptNotificationSchema = lazySchema(() =>
const ClaudeInChromePromptNotificationSchema: () => AnyObjectSchema = lazySchema(() =>
z.object({
method: z.literal('notifications/message'),
params: z.object({
Expand Down Expand Up @@ -48,7 +49,7 @@ export function usePromptsFromClaudeInChrome(
}

if (mcpClient) {
mcpClient.client.setNotificationHandler(ClaudeInChromePromptNotificationSchema() as any, notification => {
mcpClient.client.setNotificationHandler(ClaudeInChromePromptNotificationSchema(), notification => {
if (mcpClientRef.current !== mcpClient) {
return;
}
Expand Down
34 changes: 18 additions & 16 deletions src/services/mcp/channelNotification.ts
Original file line number Diff line number Diff line change
Expand Up @@ -17,6 +17,7 @@
*/

import type { ServerCapabilities } from '@modelcontextprotocol/sdk/types.js'
import type { AnyObjectSchema } from '@modelcontextprotocol/sdk/server/zod-compat.js'
import { z } from 'zod/v4'
import { type ChannelEntry, getAllowedChannels } from '../../bootstrap/state.js'
import { CHANNEL_TAG } from '../../constants/xml.js'
Expand Down Expand Up @@ -96,23 +97,24 @@ export type ChannelPermissionRequestParams = {
}
}

export const ChannelPermissionRequestNotificationSchema = lazySchema(() =>
z.object({
method: z.literal(CHANNEL_PERMISSION_REQUEST_METHOD),
params: z.object({
request_id: z.string(),
tool_name: z.string(),
description: z.string(),
input_preview: z.string(),
channel_context: z
.object({
source_server: z.string().optional(),
chat_id: z.string().optional(),
})
.optional(),
export const ChannelPermissionRequestNotificationSchema: () => AnyObjectSchema =
lazySchema(() =>
z.object({
method: z.literal(CHANNEL_PERMISSION_REQUEST_METHOD),
params: z.object({
request_id: z.string(),
tool_name: z.string(),
description: z.string(),
input_preview: z.string(),
channel_context: z
.object({
source_server: z.string().optional(),
chat_id: z.string().optional(),
})
.optional(),
}),
}),
}),
)
)

/**
* Meta keys become XML attribute NAMES — a crafted key like
Expand Down
35 changes: 25 additions & 10 deletions src/utils/forkedAgent.ts
Original file line number Diff line number Diff line change
Expand Up @@ -20,10 +20,14 @@ import {
} from '../services/analytics/index.js'
import { accumulateUsage, updateUsage } from '../services/api/claude.js'
import { EMPTY_USAGE, type NonNullableUsage } from '@ant/model-provider'
import type {
BetaRawMessageDeltaEvent,
BetaRawMessageStreamEvent,
} from '@anthropic-ai/sdk/resources/beta/messages/messages.js'
import type { ToolUseContext } from '../Tool.js'
import type { AgentDefinition } from '@claude-code-best/builtin-tools/tools/AgentTool/loadAgentsDir.js'
import type { AgentId } from '../types/ids.js'
import type { Message } from '../types/message.js'
import type { Message, StreamEvent } from '../types/message.js'
import { createChildAbortController } from './abortController.js'
import { logForDebugging } from './debug.js'
import { cloneFileStateCache } from './fileStateCache.js'
Expand Down Expand Up @@ -492,6 +496,24 @@ export function createSubagentContext(
* })
* ```
*/

type StreamEventMessage = StreamEvent & {
type: 'stream_event'
event: BetaRawMessageStreamEvent
}

function isMessageDeltaStreamEvent(
message: Message | StreamEvent,
): message is StreamEventMessage & { event: BetaRawMessageDeltaEvent } {
return (
message.type === 'stream_event' &&
typeof (message as StreamEventMessage).event === 'object' &&
(message as StreamEventMessage).event !== null &&
'type' in (message as StreamEventMessage).event &&
(message as StreamEventMessage).event.type === 'message_delta'
)
}

export async function runForkedAgent({
promptMessages,
cacheSafeParams,
Expand Down Expand Up @@ -562,15 +584,8 @@ export async function runForkedAgent({
})) {
// Extract real usage from message_delta stream events (final usage per API call)
if (message.type === 'stream_event') {
if (
'event' in message &&
(message as any).event?.type === 'message_delta' &&
(message as any).event.usage
) {
const turnUsage = updateUsage(
{ ...EMPTY_USAGE },
(message as any).event.usage,
)
if (isMessageDeltaStreamEvent(message)) {
const turnUsage = updateUsage({ ...EMPTY_USAGE }, message.event.usage)
totalUsage = accumulateUsage(totalUsage, turnUsage)
}
continue
Expand Down
34 changes: 26 additions & 8 deletions src/utils/hooks/execAgentHook.ts
Original file line number Diff line number Diff line change
Expand Up @@ -8,7 +8,12 @@ import { type Tool, toolMatchesName } from '../../Tool.js'
import { SYNTHETIC_OUTPUT_TOOL_NAME } from '@claude-code-best/builtin-tools/tools/SyntheticOutputTool/SyntheticOutputTool.js'
import { ALL_AGENT_DISALLOWED_TOOLS } from '../../tools.js'
import { asAgentId } from '../../types/ids.js'
import type { Message } from '../../types/message.js'
import type {
AttachmentMessage,
Message,
RequestStartEvent,
StreamEvent,
} from '../../types/message.js'
import { createAbortController } from '../abortController.js'
import { createAttachmentMessage } from '../attachments.js'
import { createCombinedAbortSignal } from '../combinedAbortSignal.js'
Expand All @@ -30,6 +35,24 @@ import {
} from './hookHelpers.js'
import { clearSessionHooks } from './sessionHooks.js'

type QueryMessage = Message | StreamEvent | RequestStartEvent

type StructuredOutputAttachment = {
type: 'structured_output'
data: unknown
[key: string]: unknown
}

type StructuredOutputAttachmentMessage =
AttachmentMessage<StructuredOutputAttachment>

function isStructuredOutputAttachmentMessage(
message: QueryMessage,
): message is StructuredOutputAttachmentMessage {
if (message.type !== 'attachment') return false
return (message as Message).attachment?.type === 'structured_output'
}

/**
* Execute an agent-based hook using a multi-turn LLM query
*/
Expand Down Expand Up @@ -209,13 +232,8 @@ When done, return your result using the ${SYNTHETIC_OUTPUT_TOOL_NAME} tool with:
}

// Check for structured output in attachments
if (
message.type === 'attachment' &&
(message as any).attachment.type === 'structured_output'
) {
const parsed = hookResponseSchema().safeParse(
(message as any).attachment.data,
)
if (isStructuredOutputAttachmentMessage(message)) {
const parsed = hookResponseSchema().safeParse(message.attachment.data)
if (parsed.success) {
structuredOutputResult = parsed.data
logForDebugging(
Expand Down