Skip to content

Commit 2bdd279

Browse files
authored
fix: propagate abort signal to inline read tool (#21584)
1 parent 51535d8 commit 2bdd279

File tree

4 files changed

+185
-79
lines changed

4 files changed

+185
-79
lines changed

packages/opencode/src/session/prompt.ts

Lines changed: 41 additions & 40 deletions
Original file line numberDiff line numberDiff line change
@@ -559,7 +559,7 @@ NOTE: At any point in time through this workflow you should feel free to ask the
559559
}) {
560560
const { task, model, lastUser, sessionID, session, msgs } = input
561561
const ctx = yield* InstanceState.context
562-
const taskTool = yield* registry.fromID(TaskTool.id)
562+
const { task: taskTool } = yield* registry.named()
563563
const taskModel = task.model ? yield* getModel(task.model.providerID, task.model.modelID, sessionID) : model
564564
const assistantMessage: MessageV2.Assistant = yield* sessions.updateMessage({
565565
id: MessageID.ascending(),
@@ -1080,6 +1080,21 @@ NOTE: At any point in time through this workflow you should feel free to ask the
10801080
const filepath = fileURLToPath(part.url)
10811081
if (yield* fsys.isDir(filepath)) part.mime = "application/x-directory"
10821082

1083+
const { read } = yield* registry.named()
1084+
const execRead = (args: Parameters<typeof read.execute>[0], extra?: Tool.Context["extra"]) =>
1085+
Effect.promise((signal: AbortSignal) =>
1086+
read.execute(args, {
1087+
sessionID: input.sessionID,
1088+
abort: signal,
1089+
agent: input.agent!,
1090+
messageID: info.id,
1091+
extra: { bypassCwdCheck: true, ...extra },
1092+
messages: [],
1093+
metadata: async () => {},
1094+
ask: async () => {},
1095+
}),
1096+
)
1097+
10831098
if (part.mime === "text/plain") {
10841099
let offset: number | undefined
10851100
let limit: number | undefined
@@ -1116,29 +1131,12 @@ NOTE: At any point in time through this workflow you should feel free to ask the
11161131
text: `Called the Read tool with the following input: ${JSON.stringify(args)}`,
11171132
},
11181133
]
1119-
const read = yield* registry.fromID("read").pipe(
1120-
Effect.flatMap((t) =>
1121-
provider.getModel(info.model.providerID, info.model.modelID).pipe(
1122-
Effect.flatMap((mdl) =>
1123-
Effect.promise(() =>
1124-
t.execute(args, {
1125-
sessionID: input.sessionID,
1126-
abort: new AbortController().signal,
1127-
agent: input.agent!,
1128-
messageID: info.id,
1129-
extra: { bypassCwdCheck: true, model: mdl },
1130-
messages: [],
1131-
metadata: async () => {},
1132-
ask: async () => {},
1133-
}),
1134-
),
1135-
),
1136-
),
1137-
),
1134+
const exit = yield* provider.getModel(info.model.providerID, info.model.modelID).pipe(
1135+
Effect.flatMap((mdl) => execRead(args, { model: mdl })),
11381136
Effect.exit,
11391137
)
1140-
if (Exit.isSuccess(read)) {
1141-
const result = read.value
1138+
if (Exit.isSuccess(exit)) {
1139+
const result = exit.value
11421140
pieces.push({
11431141
messageID: info.id,
11441142
sessionID: input.sessionID,
@@ -1160,7 +1158,7 @@ NOTE: At any point in time through this workflow you should feel free to ask the
11601158
pieces.push({ ...part, messageID: info.id, sessionID: input.sessionID })
11611159
}
11621160
} else {
1163-
const error = Cause.squash(read.cause)
1161+
const error = Cause.squash(exit.cause)
11641162
log.error("failed to read file", { error })
11651163
const message = error instanceof Error ? error.message : String(error)
11661164
yield* bus.publish(Session.Event.Error, {
@@ -1180,22 +1178,25 @@ NOTE: At any point in time through this workflow you should feel free to ask the
11801178

11811179
if (part.mime === "application/x-directory") {
11821180
const args = { filePath: filepath }
1183-
const result = yield* registry.fromID("read").pipe(
1184-
Effect.flatMap((t) =>
1185-
Effect.promise(() =>
1186-
t.execute(args, {
1187-
sessionID: input.sessionID,
1188-
abort: new AbortController().signal,
1189-
agent: input.agent!,
1190-
messageID: info.id,
1191-
extra: { bypassCwdCheck: true },
1192-
messages: [],
1193-
metadata: async () => {},
1194-
ask: async () => {},
1195-
}),
1196-
),
1197-
),
1198-
)
1181+
const exit = yield* execRead(args).pipe(Effect.exit)
1182+
if (Exit.isFailure(exit)) {
1183+
const error = Cause.squash(exit.cause)
1184+
log.error("failed to read directory", { error })
1185+
const message = error instanceof Error ? error.message : String(error)
1186+
yield* bus.publish(Session.Event.Error, {
1187+
sessionID: input.sessionID,
1188+
error: new NamedError.Unknown({ message }).toObject(),
1189+
})
1190+
return [
1191+
{
1192+
messageID: info.id,
1193+
sessionID: input.sessionID,
1194+
type: "text",
1195+
synthetic: true,
1196+
text: `Read tool failed to read ${filepath} with the following error: ${message}`,
1197+
},
1198+
]
1199+
}
11991200
return [
12001201
{
12011202
messageID: info.id,
@@ -1209,7 +1210,7 @@ NOTE: At any point in time through this workflow you should feel free to ask the
12091210
sessionID: input.sessionID,
12101211
type: "text",
12111212
synthetic: true,
1212-
text: result.output,
1213+
text: exit.value.output,
12131214
},
12141215
{ ...part, messageID: info.id, sessionID: input.sessionID },
12151216
]

packages/opencode/src/tool/registry.ts

Lines changed: 14 additions & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -42,24 +42,25 @@ import { Agent } from "../agent/agent"
4242
export namespace ToolRegistry {
4343
const log = Log.create({ service: "tool.registry" })
4444

45+
type TaskDef = Tool.InferDef<typeof TaskTool>
46+
type ReadDef = Tool.InferDef<typeof ReadTool>
47+
4548
type State = {
4649
custom: Tool.Def[]
4750
builtin: Tool.Def[]
51+
task: TaskDef
52+
read: ReadDef
4853
}
4954

5055
export interface Interface {
5156
readonly ids: () => Effect.Effect<string[]>
5257
readonly all: () => Effect.Effect<Tool.Def[]>
53-
readonly named: {
54-
task: Tool.Info
55-
read: Tool.Info
56-
}
58+
readonly named: () => Effect.Effect<{ task: TaskDef; read: ReadDef }>
5759
readonly tools: (model: {
5860
providerID: ProviderID
5961
modelID: ModelID
6062
agent: Agent.Info
6163
}) => Effect.Effect<Tool.Def[]>
62-
readonly fromID: (id: string) => Effect.Effect<Tool.Def>
6364
}
6465

6566
export class Service extends ServiceMap.Service<Service, Interface>()("@opencode/ToolRegistry") {}
@@ -183,6 +184,8 @@ export namespace ToolRegistry {
183184
...(Flag.OPENCODE_EXPERIMENTAL_LSP_TOOL ? [tool.lsp] : []),
184185
...(Flag.OPENCODE_EXPERIMENTAL_PLAN_MODE && Flag.OPENCODE_CLIENT === "cli" ? [tool.plan] : []),
185186
],
187+
task: tool.task,
188+
read: tool.read,
186189
}
187190
}),
188191
)
@@ -192,13 +195,6 @@ export namespace ToolRegistry {
192195
return [...s.builtin, ...s.custom] as Tool.Def[]
193196
})
194197

195-
const fromID: Interface["fromID"] = Effect.fn("ToolRegistry.fromID")(function* (id: string) {
196-
const tools = yield* all()
197-
const match = tools.find((tool) => tool.id === id)
198-
if (!match) return yield* Effect.die(`Tool not found: ${id}`)
199-
return match
200-
})
201-
202198
const ids: Interface["ids"] = Effect.fn("ToolRegistry.ids")(function* () {
203199
return (yield* all()).map((tool) => tool.id)
204200
})
@@ -245,7 +241,12 @@ export namespace ToolRegistry {
245241
)
246242
})
247243

248-
return Service.of({ ids, all, named: { task, read }, tools, fromID })
244+
const named: Interface["named"] = Effect.fn("ToolRegistry.named")(function* () {
245+
const s = yield* InstanceState.get(state)
246+
return { task: s.task, read: s.read }
247+
})
248+
249+
return Service.of({ ids, all, named, tools })
249250
}),
250251
)
251252

packages/opencode/src/tool/tool.ts

Lines changed: 8 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -60,6 +60,13 @@ export namespace Tool {
6060
export type InferMetadata<T> =
6161
T extends Info<any, infer M> ? M : T extends Effect.Effect<Info<any, infer M>, any, any> ? M : never
6262

63+
export type InferDef<T> =
64+
T extends Info<infer P, infer M>
65+
? Def<P, M>
66+
: T extends Effect.Effect<Info<infer P, infer M>, any, any>
67+
? Def<P, M>
68+
: never
69+
6370
function wrap<Parameters extends z.ZodType, Result extends Metadata>(
6471
id: string,
6572
init: (() => Promise<DefWithoutID<Parameters, Result>>) | DefWithoutID<Parameters, Result>,
@@ -118,7 +125,7 @@ export namespace Tool {
118125
)
119126
}
120127

121-
export function init(info: Info): Effect.Effect<Def> {
128+
export function init<P extends z.ZodType, M extends Metadata>(info: Info<P, M>): Effect.Effect<Def<P, M>> {
122129
return Effect.gen(function* () {
123130
const init = yield* Effect.promise(() => info.init())
124131
return {

packages/opencode/test/session/prompt-effect.test.ts

Lines changed: 122 additions & 25 deletions
Original file line numberDiff line numberDiff line change
@@ -631,31 +631,22 @@ it.live(
631631
const ready = defer<void>()
632632
const aborted = defer<void>()
633633
const registry = yield* ToolRegistry.Service
634-
const init = registry.named.task.init
635-
registry.named.task.init = async () => ({
636-
description: "task",
637-
parameters: z.object({
638-
description: z.string(),
639-
prompt: z.string(),
640-
subagent_type: z.string(),
641-
task_id: z.string().optional(),
642-
command: z.string().optional(),
643-
}),
644-
execute: async (_args, ctx) => {
645-
ready.resolve()
646-
ctx.abort.addEventListener("abort", () => aborted.resolve(), { once: true })
647-
await new Promise<void>(() => {})
648-
return {
649-
title: "",
650-
metadata: {
651-
sessionId: SessionID.make("task"),
652-
model: ref,
653-
},
654-
output: "",
655-
}
656-
},
657-
})
658-
yield* Effect.addFinalizer(() => Effect.sync(() => void (registry.named.task.init = init)))
634+
const { task } = yield* registry.named()
635+
const original = task.execute
636+
task.execute = async (_args, ctx) => {
637+
ready.resolve()
638+
ctx.abort.addEventListener("abort", () => aborted.resolve(), { once: true })
639+
await new Promise<void>(() => {})
640+
return {
641+
title: "",
642+
metadata: {
643+
sessionId: SessionID.make("task"),
644+
model: ref,
645+
},
646+
output: "",
647+
}
648+
}
649+
yield* Effect.addFinalizer(() => Effect.sync(() => void (task.execute = original)))
659650

660651
const { prompt, chat } = yield* boot()
661652
const msg = yield* user(chat.id, "hello")
@@ -1240,3 +1231,109 @@ unix(
12401231
),
12411232
30_000,
12421233
)
1234+
1235+
// Abort signal propagation tests for inline tool execution
1236+
1237+
/** Override a tool's execute to hang until aborted. Returns ready/aborted defers and a finalizer. */
1238+
function hangUntilAborted(tool: { execute: (...args: any[]) => any }) {
1239+
const ready = defer<void>()
1240+
const aborted = defer<void>()
1241+
const original = tool.execute
1242+
tool.execute = async (_args: any, ctx: any) => {
1243+
ready.resolve()
1244+
ctx.abort.addEventListener("abort", () => aborted.resolve(), { once: true })
1245+
await new Promise<void>(() => {})
1246+
return { title: "", metadata: {}, output: "" }
1247+
}
1248+
const restore = Effect.addFinalizer(() => Effect.sync(() => void (tool.execute = original)))
1249+
return { ready, aborted, restore }
1250+
}
1251+
1252+
it.live(
1253+
"interrupt propagates abort signal to read tool via file part (text/plain)",
1254+
() =>
1255+
provideTmpdirInstance(
1256+
(dir) =>
1257+
Effect.gen(function* () {
1258+
const registry = yield* ToolRegistry.Service
1259+
const { read } = yield* registry.named()
1260+
const { ready, aborted, restore } = hangUntilAborted(read)
1261+
yield* restore
1262+
1263+
const prompt = yield* SessionPrompt.Service
1264+
const sessions = yield* Session.Service
1265+
const chat = yield* sessions.create({ title: "Abort Test" })
1266+
1267+
const testFile = path.join(dir, "test.txt")
1268+
yield* Effect.promise(() => Bun.write(testFile, "hello world"))
1269+
1270+
const fiber = yield* prompt
1271+
.prompt({
1272+
sessionID: chat.id,
1273+
agent: "build",
1274+
parts: [
1275+
{ type: "text", text: "read this" },
1276+
{ type: "file", url: `file://${testFile}`, filename: "test.txt", mime: "text/plain" },
1277+
],
1278+
})
1279+
.pipe(Effect.forkChild)
1280+
1281+
yield* Effect.promise(() => ready.promise)
1282+
yield* Fiber.interrupt(fiber)
1283+
1284+
yield* Effect.promise(() =>
1285+
Promise.race([
1286+
aborted.promise,
1287+
new Promise<void>((_, reject) =>
1288+
setTimeout(() => reject(new Error("abort signal not propagated within 2s")), 2_000),
1289+
),
1290+
]),
1291+
)
1292+
}),
1293+
{ git: true, config: cfg },
1294+
),
1295+
30_000,
1296+
)
1297+
1298+
it.live(
1299+
"interrupt propagates abort signal to read tool via file part (directory)",
1300+
() =>
1301+
provideTmpdirInstance(
1302+
(dir) =>
1303+
Effect.gen(function* () {
1304+
const registry = yield* ToolRegistry.Service
1305+
const { read } = yield* registry.named()
1306+
const { ready, aborted, restore } = hangUntilAborted(read)
1307+
yield* restore
1308+
1309+
const prompt = yield* SessionPrompt.Service
1310+
const sessions = yield* Session.Service
1311+
const chat = yield* sessions.create({ title: "Abort Test" })
1312+
1313+
const fiber = yield* prompt
1314+
.prompt({
1315+
sessionID: chat.id,
1316+
agent: "build",
1317+
parts: [
1318+
{ type: "text", text: "read this" },
1319+
{ type: "file", url: `file://${dir}`, filename: "dir", mime: "application/x-directory" },
1320+
],
1321+
})
1322+
.pipe(Effect.forkChild)
1323+
1324+
yield* Effect.promise(() => ready.promise)
1325+
yield* Fiber.interrupt(fiber)
1326+
1327+
yield* Effect.promise(() =>
1328+
Promise.race([
1329+
aborted.promise,
1330+
new Promise<void>((_, reject) =>
1331+
setTimeout(() => reject(new Error("abort signal not propagated within 2s")), 2_000),
1332+
),
1333+
]),
1334+
)
1335+
}),
1336+
{ git: true, config: cfg },
1337+
),
1338+
30_000,
1339+
)

0 commit comments

Comments
 (0)