feat(core): support background shell tool
This commit is contained in:
Dax Raad
parent
ff4cab03c1
commit
19a5b5a05d
6 changed files with 344 additions and 193 deletions
|
|
@ -2,26 +2,38 @@ import fs from "fs/promises"
|
|||
import { realpathSync } from "node:fs"
|
||||
import path from "path"
|
||||
import { describe, expect, test } from "bun:test"
|
||||
import { Effect, Layer } from "effect"
|
||||
import { Config } from "@opencode-ai/core/config"
|
||||
import { DateTime, Effect, Layer } from "effect"
|
||||
import { AppNodeBuilder } from "@opencode-ai/core/effect/app-node-builder"
|
||||
import { LayerNode } from "@opencode-ai/core/effect/layer-node"
|
||||
import { makeGlobalNode } from "@opencode-ai/core/effect/app-node"
|
||||
import { filesystem } from "@opencode-ai/core/effect/app-node-platform"
|
||||
import { Database } from "@opencode-ai/core/database/database"
|
||||
import { EventV2 } from "@opencode-ai/core/event"
|
||||
import { FSUtil } from "@opencode-ai/core/fs-util"
|
||||
import { Global } from "@opencode-ai/core/global"
|
||||
import { Location } from "@opencode-ai/core/location"
|
||||
import { LocationMutation } from "@opencode-ai/core/location-mutation"
|
||||
import { PermissionV2 } from "@opencode-ai/core/permission"
|
||||
import { AppProcess } from "@opencode-ai/core/process"
|
||||
import { Project } from "@opencode-ai/core/project"
|
||||
import { LocationServiceMap } from "@opencode-ai/core/location-service-map"
|
||||
import { ModelV2 } from "@opencode-ai/core/model"
|
||||
import { ProviderV2 } from "@opencode-ai/core/provider"
|
||||
import { AbsolutePath } from "@opencode-ai/core/schema"
|
||||
import { AgentV2 } from "@opencode-ai/core/agent"
|
||||
import { BackgroundJob } from "@opencode-ai/core/background-job"
|
||||
import { SessionV2 } from "@opencode-ai/core/session"
|
||||
import { Shell } from "@opencode-ai/core/shell"
|
||||
import { SessionEvent } from "@opencode-ai/core/session/event"
|
||||
import { SessionExecution } from "@opencode-ai/core/session/execution"
|
||||
import { SessionMessage } from "@opencode-ai/core/session/message"
|
||||
import { SessionRunnerModel } from "@opencode-ai/core/session/runner/model"
|
||||
import { SessionStore } from "@opencode-ai/core/session/store"
|
||||
import { PermissionV2 } from "@opencode-ai/core/permission"
|
||||
import { ShellTool } from "@opencode-ai/core/tool/shell"
|
||||
import { ToolRegistry } from "@opencode-ai/core/tool/registry"
|
||||
import { ToolOutputStore } from "@opencode-ai/core/tool-output-store"
|
||||
import { tmpdir } from "./fixture/tmpdir"
|
||||
import { testEffect } from "./lib/effect"
|
||||
import { toolIdentity, executeTool, settleTool, toolDefinitions } from "./lib/tool"
|
||||
|
||||
const sessionID = SessionV2.ID.make("ses_shell_tool_test")
|
||||
const sessionModel = ModelV2.Ref.make({ id: ModelV2.ID.make("test"), providerID: ProviderV2.ID.make("test") })
|
||||
const assertions: PermissionV2.AssertInput[] = []
|
||||
let denyAction: string | undefined
|
||||
let afterPermission = (_input: PermissionV2.AssertInput): Effect.Effect<void> => Effect.void
|
||||
|
|
@ -50,37 +62,80 @@ const reset = () => {
|
|||
afterPermission = () => Effect.void
|
||||
}
|
||||
|
||||
const withTool = <A, E, R>(
|
||||
data: string,
|
||||
directory: string,
|
||||
body: (registry: ToolRegistry.Interface) => Effect.Effect<A, E, R>,
|
||||
) => {
|
||||
const filesystem = FSUtil.defaultLayer
|
||||
const location = Location.layer(Location.Ref.make({ directory: AbsolutePath.make(directory) })).pipe(
|
||||
Layer.provide(Project.defaultLayer),
|
||||
)
|
||||
const global = Global.layerWith({ data, config: path.join(data, "config") })
|
||||
const mutation = LocationMutation.layer.pipe(Layer.provide(filesystem), Layer.provide(location))
|
||||
const registry = ToolRegistry.defaultLayer.pipe(Layer.provide(permission))
|
||||
const shellService = Shell.layer.pipe(
|
||||
Layer.provide(EventV2.defaultLayer),
|
||||
Layer.provide(location),
|
||||
Layer.provide(Config.locationLayer.pipe(Layer.provide(location), Layer.provide(filesystem), Layer.provide(global))),
|
||||
Layer.provide(global),
|
||||
Layer.provide(filesystem),
|
||||
Layer.provide(AppProcess.defaultLayer),
|
||||
)
|
||||
const shell = ShellTool.layer.pipe(
|
||||
Layer.provide(registry),
|
||||
Layer.provide(permission),
|
||||
Layer.provide(mutation),
|
||||
Layer.provide(filesystem),
|
||||
Layer.provide(shellService),
|
||||
)
|
||||
return Effect.gen(function* () {
|
||||
return yield* body(yield* ToolRegistry.Service)
|
||||
}).pipe(Effect.provide(Layer.mergeAll(registry, shell, filesystem)))
|
||||
}
|
||||
const executionNode = makeGlobalNode({
|
||||
service: SessionExecution.Service,
|
||||
layer: Layer.effect(
|
||||
SessionExecution.Service,
|
||||
Effect.gen(function* () {
|
||||
const events = yield* EventV2.Service
|
||||
const store = yield* SessionStore.Service
|
||||
const complete = Effect.fn("ShellTest.complete")(function* (id: SessionV2.ID) {
|
||||
const session = yield* store.get(id)
|
||||
if (!session) return
|
||||
const assistantMessageID = SessionMessage.ID.create()
|
||||
const textID = "text_shell_test"
|
||||
yield* events.publish(SessionEvent.Step.Started, {
|
||||
sessionID: id,
|
||||
assistantMessageID,
|
||||
timestamp: yield* DateTime.now,
|
||||
agent: session.agent ?? AgentV2.ID.make("code"),
|
||||
model: sessionModel,
|
||||
})
|
||||
yield* events.publish(SessionEvent.Text.Started, {
|
||||
sessionID: id,
|
||||
assistantMessageID,
|
||||
timestamp: yield* DateTime.now,
|
||||
textID,
|
||||
})
|
||||
yield* events.publish(SessionEvent.Text.Ended, {
|
||||
sessionID: id,
|
||||
assistantMessageID,
|
||||
timestamp: yield* DateTime.now,
|
||||
textID,
|
||||
text: "ok",
|
||||
})
|
||||
yield* events.publish(SessionEvent.Step.Ended, {
|
||||
sessionID: id,
|
||||
assistantMessageID,
|
||||
timestamp: yield* DateTime.now,
|
||||
finish: "stop",
|
||||
cost: 0,
|
||||
tokens: { input: 0, output: 0, reasoning: 0, cache: { read: 0, write: 0 } },
|
||||
})
|
||||
})
|
||||
return SessionExecution.Service.of({
|
||||
active: Effect.succeed(new Set()),
|
||||
resume: complete,
|
||||
wake: () => Effect.void,
|
||||
interrupt: () => Effect.void,
|
||||
awaitIdle: (id) => complete(id).pipe(Effect.exit, Effect.asVoid),
|
||||
})
|
||||
}),
|
||||
),
|
||||
deps: [EventV2.node, SessionStore.node],
|
||||
})
|
||||
|
||||
const layer = AppNodeBuilder.build(
|
||||
LayerNode.bind(
|
||||
LayerNode.group([
|
||||
Database.node,
|
||||
EventV2.node,
|
||||
BackgroundJob.node,
|
||||
ToolOutputStore.cleanupNode,
|
||||
SessionV2.node,
|
||||
ShellTool.node,
|
||||
LocationServiceMap.node,
|
||||
filesystem,
|
||||
FSUtil.node,
|
||||
Global.node,
|
||||
]),
|
||||
SessionExecution.node,
|
||||
executionNode,
|
||||
),
|
||||
[LayerNode.replace(PermissionV2.layer, permission)],
|
||||
)
|
||||
|
||||
const it = testEffect(layer)
|
||||
|
||||
const call = (input: typeof ShellTool.Input.Type, id = "call-shell") => ({
|
||||
sessionID,
|
||||
|
|
@ -100,20 +155,42 @@ const overflowCommand = (bytes: number) =>
|
|||
? `[Console]::Out.Write(('x' * ${bytes})); Start-Sleep -Milliseconds 100`
|
||||
: `head -c ${bytes} /dev/zero | tr '\\0' 'x'`
|
||||
|
||||
const it = testEffect(Layer.empty)
|
||||
const withSession = <A, E, R>(
|
||||
directory: string,
|
||||
body: (registry: ToolRegistry.Interface) => Effect.Effect<A, E, R>,
|
||||
) =>
|
||||
Effect.gen(function* () {
|
||||
const sessions = yield* SessionV2.Service
|
||||
const location = Location.Ref.make({ directory: AbsolutePath.make(directory) })
|
||||
yield* sessions.create({
|
||||
id: sessionID,
|
||||
title: "shell test",
|
||||
location,
|
||||
model: sessionModel,
|
||||
})
|
||||
const locations = yield* LocationServiceMap.Service
|
||||
const locationLayer = locations.get(location)
|
||||
const registry = yield* ToolRegistry.Service.pipe(Effect.provide(locationLayer))
|
||||
return yield* body(registry).pipe(Effect.provide(locationLayer))
|
||||
})
|
||||
|
||||
describe("ShellTool", () => {
|
||||
it.live("registers and returns real successful output from the active Location", () =>
|
||||
Effect.acquireUseRelease(
|
||||
Effect.promise(() => Promise.all([tmpdir(), tmpdir()])),
|
||||
([data, tmp]) => {
|
||||
Effect.promise(() => tmpdir()),
|
||||
(tmp) => {
|
||||
reset()
|
||||
return withTool(data.path, tmp.path, (registry) =>
|
||||
return withSession(tmp.path, (registry) =>
|
||||
Effect.gen(function* () {
|
||||
const definitions = yield* toolDefinitions(registry)
|
||||
expect(definitions.map((tool) => tool.name)).toEqual(["shell"])
|
||||
expect(definitions[0]?.outputSchema).not.toHaveProperty("properties.output")
|
||||
expect(yield* toolDefinitions(registry, [{ action: "shell", resource: "*", effect: "deny" }])).toEqual([])
|
||||
const shell = definitions.find((tool) => tool.name === "shell")
|
||||
expect(shell).toBeDefined()
|
||||
expect(shell?.outputSchema).not.toHaveProperty("properties.output")
|
||||
expect(
|
||||
(yield* toolDefinitions(registry, [{ action: "shell", resource: "*", effect: "deny" }])).map(
|
||||
(tool) => tool.name,
|
||||
),
|
||||
).not.toContain("shell")
|
||||
|
||||
const settled = yield* settleTool(registry, call({ command: helloCommand }))
|
||||
expect(settled.output?.structured).toMatchObject({ exit: 0, truncated: false })
|
||||
|
|
@ -126,21 +203,18 @@ describe("ShellTool", () => {
|
|||
}),
|
||||
)
|
||||
},
|
||||
([data, tmp]) =>
|
||||
Effect.promise(() =>
|
||||
Promise.all([data[Symbol.asyncDispose](), tmp[Symbol.asyncDispose]()]).then(() => undefined),
|
||||
),
|
||||
(tmp) => Effect.promise(() => tmp[Symbol.asyncDispose]().then(() => undefined)),
|
||||
),
|
||||
)
|
||||
|
||||
it.live("resolves a relative workdir from the active Location", () =>
|
||||
Effect.acquireUseRelease(
|
||||
Effect.promise(() => Promise.all([tmpdir(), tmpdir()])),
|
||||
([data, tmp]) => {
|
||||
Effect.promise(() => tmpdir()),
|
||||
(tmp) => {
|
||||
reset()
|
||||
return Effect.promise(() => fs.mkdir(path.join(tmp.path, "src"))).pipe(
|
||||
Effect.andThen(
|
||||
withTool(data.path, tmp.path, (registry) =>
|
||||
withSession(tmp.path, (registry) =>
|
||||
settleTool(registry, call({ command: cwdCommand, workdir: "src" })),
|
||||
),
|
||||
),
|
||||
|
|
@ -154,17 +228,14 @@ describe("ShellTool", () => {
|
|||
),
|
||||
)
|
||||
},
|
||||
([data, tmp]) =>
|
||||
Effect.promise(() =>
|
||||
Promise.all([data[Symbol.asyncDispose](), tmp[Symbol.asyncDispose]()]).then(() => undefined),
|
||||
),
|
||||
(tmp) => Effect.promise(() => tmp[Symbol.asyncDispose]().then(() => undefined)),
|
||||
),
|
||||
)
|
||||
|
||||
it.live("rejects a workdir that stops being a directory during approval", () =>
|
||||
Effect.acquireUseRelease(
|
||||
Effect.promise(() => Promise.all([tmpdir(), tmpdir()])),
|
||||
([data, tmp]) => {
|
||||
Effect.promise(() => tmpdir()),
|
||||
(tmp) => {
|
||||
reset()
|
||||
const workdir = path.join(tmp.path, "src")
|
||||
afterPermission = (input) =>
|
||||
|
|
@ -176,26 +247,23 @@ describe("ShellTool", () => {
|
|||
: Effect.void
|
||||
return Effect.promise(() => fs.mkdir(workdir)).pipe(
|
||||
Effect.andThen(
|
||||
withTool(data.path, tmp.path, (registry) =>
|
||||
withSession(tmp.path, (registry) =>
|
||||
executeTool(registry, call({ command: cwdCommand, workdir: "src" })),
|
||||
),
|
||||
),
|
||||
Effect.andThen(Effect.sync(() => expect(assertions.map((input) => input.action)).toEqual(["shell"]))),
|
||||
)
|
||||
},
|
||||
([data, tmp]) =>
|
||||
Effect.promise(() =>
|
||||
Promise.all([data[Symbol.asyncDispose](), tmp[Symbol.asyncDispose]()]).then(() => undefined),
|
||||
),
|
||||
(tmp) => Effect.promise(() => tmp[Symbol.asyncDispose]().then(() => undefined)),
|
||||
),
|
||||
)
|
||||
|
||||
it.live("approves an explicit external workdir before shell execution", () =>
|
||||
Effect.acquireUseRelease(
|
||||
Effect.promise(() => Promise.all([tmpdir(), tmpdir(), tmpdir()])),
|
||||
([data, active, outside]) => {
|
||||
Effect.promise(() => Promise.all([tmpdir(), tmpdir()])),
|
||||
([active, outside]) => {
|
||||
reset()
|
||||
return withTool(data.path, active.path, (registry) =>
|
||||
return withSession(active.path, (registry) =>
|
||||
executeTool(registry, call({ command: cwdCommand, workdir: outside.path })),
|
||||
).pipe(
|
||||
Effect.andThen(
|
||||
|
|
@ -208,53 +276,45 @@ describe("ShellTool", () => {
|
|||
),
|
||||
)
|
||||
},
|
||||
([data, active, outside]) =>
|
||||
([active, outside]) =>
|
||||
Effect.promise(() =>
|
||||
Promise.all([
|
||||
data[Symbol.asyncDispose](),
|
||||
active[Symbol.asyncDispose](),
|
||||
outside[Symbol.asyncDispose](),
|
||||
]).then(() => undefined),
|
||||
Promise.all([active[Symbol.asyncDispose](), outside[Symbol.asyncDispose]()]).then(() => undefined),
|
||||
),
|
||||
),
|
||||
)
|
||||
|
||||
it.live("does not execute after external-directory or shell denial", () =>
|
||||
Effect.acquireUseRelease(
|
||||
Effect.promise(() => Promise.all([tmpdir(), tmpdir(), tmpdir()])),
|
||||
([data, active, outside]) =>
|
||||
Effect.promise(() => Promise.all([tmpdir(), tmpdir()])),
|
||||
([active, outside]) =>
|
||||
Effect.gen(function* () {
|
||||
reset()
|
||||
denyAction = "external_directory"
|
||||
yield* withTool(data.path, active.path, (registry) =>
|
||||
yield* withSession(active.path, (registry) =>
|
||||
executeTool(registry, call({ command: cwdCommand, workdir: outside.path })),
|
||||
)
|
||||
expect(assertions.map((item) => item.action)).toEqual(["external_directory"])
|
||||
|
||||
reset()
|
||||
denyAction = "shell"
|
||||
yield* withTool(data.path, active.path, (registry) => executeTool(registry, call({ command: cwdCommand })))
|
||||
yield* withSession(active.path, (registry) => executeTool(registry, call({ command: cwdCommand })))
|
||||
expect(assertions.map((item) => item.action)).toEqual(["shell"])
|
||||
}),
|
||||
([data, active, outside]) =>
|
||||
([active, outside]) =>
|
||||
Effect.promise(() =>
|
||||
Promise.all([
|
||||
data[Symbol.asyncDispose](),
|
||||
active[Symbol.asyncDispose](),
|
||||
outside[Symbol.asyncDispose](),
|
||||
]).then(() => undefined),
|
||||
Promise.all([active[Symbol.asyncDispose](), outside[Symbol.asyncDispose]()]).then(() => undefined),
|
||||
),
|
||||
),
|
||||
)
|
||||
|
||||
it.live("reports external command arguments as advisory warnings without enforcing approval", () =>
|
||||
Effect.acquireUseRelease(
|
||||
Effect.promise(() => Promise.all([tmpdir(), tmpdir(), tmpdir()])),
|
||||
([data, active, outside]) => {
|
||||
Effect.promise(() => Promise.all([tmpdir(), tmpdir()])),
|
||||
([active, outside]) => {
|
||||
reset()
|
||||
denyAction = "external_directory"
|
||||
const target = path.join(outside.path, "secret.txt")
|
||||
return withTool(data.path, active.path, (registry) =>
|
||||
return withSession(active.path, (registry) =>
|
||||
settleTool(registry, call({ command: `cat ${target}` })),
|
||||
).pipe(
|
||||
Effect.andThen((settled) =>
|
||||
|
|
@ -269,23 +329,19 @@ describe("ShellTool", () => {
|
|||
),
|
||||
)
|
||||
},
|
||||
([data, active, outside]) =>
|
||||
([active, outside]) =>
|
||||
Effect.promise(() =>
|
||||
Promise.all([
|
||||
data[Symbol.asyncDispose](),
|
||||
active[Symbol.asyncDispose](),
|
||||
outside[Symbol.asyncDispose](),
|
||||
]).then(() => undefined),
|
||||
Promise.all([active[Symbol.asyncDispose](), outside[Symbol.asyncDispose]()]).then(() => undefined),
|
||||
),
|
||||
),
|
||||
)
|
||||
|
||||
it.live("keeps non-zero exits useful", () =>
|
||||
Effect.acquireUseRelease(
|
||||
Effect.promise(() => Promise.all([tmpdir(), tmpdir()])),
|
||||
([data, tmp]) => {
|
||||
Effect.promise(() => tmpdir()),
|
||||
(tmp) => {
|
||||
reset()
|
||||
return withTool(data.path, tmp.path, (registry) =>
|
||||
return withSession(tmp.path, (registry) =>
|
||||
settleTool(registry, call({ command: bodyExitCommand }, "call-nonzero")),
|
||||
).pipe(
|
||||
Effect.andThen((settled) =>
|
||||
|
|
@ -300,20 +356,17 @@ describe("ShellTool", () => {
|
|||
),
|
||||
)
|
||||
},
|
||||
([data, tmp]) =>
|
||||
Effect.promise(() =>
|
||||
Promise.all([data[Symbol.asyncDispose](), tmp[Symbol.asyncDispose]()]).then(() => undefined),
|
||||
),
|
||||
(tmp) => Effect.promise(() => tmp[Symbol.asyncDispose]().then(() => undefined)),
|
||||
),
|
||||
)
|
||||
|
||||
it.live("truncates the model view and points at the saved output file when output overflows", () =>
|
||||
Effect.acquireUseRelease(
|
||||
Effect.promise(() => Promise.all([tmpdir(), tmpdir()])),
|
||||
([data, tmp]) => {
|
||||
Effect.promise(() => tmpdir()),
|
||||
(tmp) => {
|
||||
reset()
|
||||
const bytes = ShellTool.MAX_CAPTURE_BYTES + 1024
|
||||
return withTool(data.path, tmp.path, (registry) =>
|
||||
return withSession(tmp.path, (registry) =>
|
||||
settleTool(registry, call({ command: overflowCommand(bytes) }, "call-overflow")),
|
||||
).pipe(
|
||||
Effect.andThen((settled) =>
|
||||
|
|
@ -327,19 +380,16 @@ describe("ShellTool", () => {
|
|||
),
|
||||
)
|
||||
},
|
||||
([data, tmp]) =>
|
||||
Effect.promise(() =>
|
||||
Promise.all([data[Symbol.asyncDispose](), tmp[Symbol.asyncDispose]()]).then(() => undefined),
|
||||
),
|
||||
(tmp) => Effect.promise(() => tmp[Symbol.asyncDispose]().then(() => undefined)),
|
||||
),
|
||||
)
|
||||
|
||||
it.live("returns a useful timeout settlement", () =>
|
||||
Effect.acquireUseRelease(
|
||||
Effect.promise(() => Promise.all([tmpdir(), tmpdir()])),
|
||||
([data, tmp]) => {
|
||||
Effect.promise(() => tmpdir()),
|
||||
(tmp) => {
|
||||
reset()
|
||||
return withTool(data.path, tmp.path, (registry) =>
|
||||
return withSession(tmp.path, (registry) =>
|
||||
settleTool(registry, call({ command: idleCommand, timeout: 50 })),
|
||||
).pipe(
|
||||
Effect.andThen((settled) =>
|
||||
|
|
@ -353,10 +403,7 @@ describe("ShellTool", () => {
|
|||
),
|
||||
)
|
||||
},
|
||||
([data, tmp]) =>
|
||||
Effect.promise(() =>
|
||||
Promise.all([data[Symbol.asyncDispose](), tmp[Symbol.asyncDispose]()]).then(() => undefined),
|
||||
),
|
||||
(tmp) => Effect.promise(() => tmp[Symbol.asyncDispose]().then(() => undefined)),
|
||||
),
|
||||
)
|
||||
})
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue