fix: restore instance context in prompt runs (#22498)
This commit is contained in:
@@ -1,5 +1,5 @@
|
|||||||
import { Layer, ManagedRuntime } from "effect"
|
import { Layer, ManagedRuntime } from "effect"
|
||||||
import { memoMap } from "./run-service"
|
import { attach, memoMap } from "./run-service"
|
||||||
import { Observability } from "./oltp"
|
import { Observability } from "./oltp"
|
||||||
|
|
||||||
import { AppFileSystem } from "@/filesystem"
|
import { AppFileSystem } from "@/filesystem"
|
||||||
@@ -97,4 +97,25 @@ export const AppLayer = Layer.mergeAll(
|
|||||||
SessionShare.defaultLayer,
|
SessionShare.defaultLayer,
|
||||||
)
|
)
|
||||||
|
|
||||||
export const AppRuntime = ManagedRuntime.make(AppLayer, { memoMap })
|
const rt = ManagedRuntime.make(AppLayer, { memoMap })
|
||||||
|
type Runtime = Pick<typeof rt, "runSync" | "runPromise" | "runPromiseExit" | "runFork" | "runCallback" | "dispose">
|
||||||
|
const wrap = (effect: Parameters<typeof rt.runSync>[0]) => attach(effect as never) as never
|
||||||
|
|
||||||
|
export const AppRuntime: Runtime = {
|
||||||
|
runSync(effect) {
|
||||||
|
return rt.runSync(wrap(effect))
|
||||||
|
},
|
||||||
|
runPromise(effect, options) {
|
||||||
|
return rt.runPromise(wrap(effect), options)
|
||||||
|
},
|
||||||
|
runPromiseExit(effect, options) {
|
||||||
|
return rt.runPromiseExit(wrap(effect), options)
|
||||||
|
},
|
||||||
|
runFork(effect) {
|
||||||
|
return rt.runFork(wrap(effect))
|
||||||
|
},
|
||||||
|
runCallback(effect) {
|
||||||
|
return rt.runCallback(wrap(effect))
|
||||||
|
},
|
||||||
|
dispose: () => rt.dispose(),
|
||||||
|
}
|
||||||
|
|||||||
@@ -104,12 +104,21 @@ export namespace SessionPrompt {
|
|||||||
const summary = yield* SessionSummary.Service
|
const summary = yield* SessionSummary.Service
|
||||||
const sys = yield* SystemPrompt.Service
|
const sys = yield* SystemPrompt.Service
|
||||||
const llm = yield* LLM.Service
|
const llm = yield* LLM.Service
|
||||||
const ctx = yield* Effect.context()
|
const runner = Effect.fn("SessionPrompt.runner")(function* () {
|
||||||
|
const ctx = yield* Effect.context()
|
||||||
const run = {
|
return {
|
||||||
promise: <A, E>(effect: Effect.Effect<A, E>) => Effect.runPromiseWith(ctx)(effect),
|
promise: <A, E>(effect: Effect.Effect<A, E>) => Effect.runPromiseWith(ctx)(effect),
|
||||||
fork: <A, E>(effect: Effect.Effect<A, E>) => Effect.runForkWith(ctx)(effect),
|
fork: <A, E>(effect: Effect.Effect<A, E>) => Effect.runForkWith(ctx)(effect),
|
||||||
}
|
}
|
||||||
|
})
|
||||||
|
const ops = Effect.fn("SessionPrompt.ops")(function* () {
|
||||||
|
const run = yield* runner()
|
||||||
|
return {
|
||||||
|
cancel: (sessionID: SessionID) => run.fork(cancel(sessionID)),
|
||||||
|
resolvePromptParts: (template: string) => resolvePromptParts(template),
|
||||||
|
prompt: (input: PromptInput) => prompt(input),
|
||||||
|
} satisfies TaskPromptOps
|
||||||
|
})
|
||||||
|
|
||||||
const cancel = Effect.fn("SessionPrompt.cancel")(function* (sessionID: SessionID) {
|
const cancel = Effect.fn("SessionPrompt.cancel")(function* (sessionID: SessionID) {
|
||||||
yield* elog.info("cancel", { sessionID })
|
yield* elog.info("cancel", { sessionID })
|
||||||
@@ -359,6 +368,8 @@ NOTE: At any point in time through this workflow you should feel free to ask the
|
|||||||
}) {
|
}) {
|
||||||
using _ = log.time("resolveTools")
|
using _ = log.time("resolveTools")
|
||||||
const tools: Record<string, AITool> = {}
|
const tools: Record<string, AITool> = {}
|
||||||
|
const run = yield* runner()
|
||||||
|
const promptOps = yield* ops()
|
||||||
|
|
||||||
const context = (args: any, options: ToolExecutionOptions): Tool.Context => ({
|
const context = (args: any, options: ToolExecutionOptions): Tool.Context => ({
|
||||||
sessionID: input.session.id,
|
sessionID: input.session.id,
|
||||||
@@ -528,6 +539,7 @@ NOTE: At any point in time through this workflow you should feel free to ask the
|
|||||||
}) {
|
}) {
|
||||||
const { task, model, lastUser, sessionID, session, msgs } = input
|
const { task, model, lastUser, sessionID, session, msgs } = input
|
||||||
const ctx = yield* InstanceState.context
|
const ctx = yield* InstanceState.context
|
||||||
|
const promptOps = yield* ops()
|
||||||
const { task: taskTool } = yield* registry.named()
|
const { task: taskTool } = yield* registry.named()
|
||||||
const taskModel = task.model ? yield* getModel(task.model.providerID, task.model.modelID, sessionID) : model
|
const taskModel = task.model ? yield* getModel(task.model.providerID, task.model.modelID, sessionID) : model
|
||||||
const assistantMessage: MessageV2.Assistant = yield* sessions.updateMessage({
|
const assistantMessage: MessageV2.Assistant = yield* sessions.updateMessage({
|
||||||
@@ -712,6 +724,7 @@ NOTE: At any point in time through this workflow you should feel free to ask the
|
|||||||
|
|
||||||
const shellImpl = Effect.fn("SessionPrompt.shellImpl")(function* (input: ShellInput) {
|
const shellImpl = Effect.fn("SessionPrompt.shellImpl")(function* (input: ShellInput) {
|
||||||
const ctx = yield* InstanceState.context
|
const ctx = yield* InstanceState.context
|
||||||
|
const run = yield* runner()
|
||||||
const session = yield* sessions.get(input.sessionID)
|
const session = yield* sessions.get(input.sessionID)
|
||||||
if (session.revert) {
|
if (session.revert) {
|
||||||
yield* revert.cleanup(session)
|
yield* revert.cleanup(session)
|
||||||
@@ -1659,12 +1672,6 @@ NOTE: At any point in time through this workflow you should feel free to ask the
|
|||||||
return result
|
return result
|
||||||
})
|
})
|
||||||
|
|
||||||
const promptOps: TaskPromptOps = {
|
|
||||||
cancel: (sessionID) => run.fork(cancel(sessionID)),
|
|
||||||
resolvePromptParts: (template) => resolvePromptParts(template),
|
|
||||||
prompt: (input) => prompt(input),
|
|
||||||
}
|
|
||||||
|
|
||||||
return Service.of({
|
return Service.of({
|
||||||
cancel,
|
cancel,
|
||||||
prompt,
|
prompt,
|
||||||
|
|||||||
@@ -1,8 +1,11 @@
|
|||||||
import { expect, test } from "bun:test"
|
import { expect, test } from "bun:test"
|
||||||
import { Context, Effect, Layer, Logger } from "effect"
|
import { Context, Effect, Layer, Logger } from "effect"
|
||||||
import { AppRuntime } from "../../src/effect/app-runtime"
|
import { AppRuntime } from "../../src/effect/app-runtime"
|
||||||
|
import { InstanceRef } from "../../src/effect/instance-ref"
|
||||||
import { EffectLogger } from "../../src/effect/logger"
|
import { EffectLogger } from "../../src/effect/logger"
|
||||||
import { makeRuntime } from "../../src/effect/run-service"
|
import { makeRuntime } from "../../src/effect/run-service"
|
||||||
|
import { Instance } from "../../src/project/instance"
|
||||||
|
import { tmpdir } from "../fixture/fixture"
|
||||||
|
|
||||||
function check(loggers: ReadonlySet<Logger.Logger<unknown, any>>) {
|
function check(loggers: ReadonlySet<Logger.Logger<unknown, any>>) {
|
||||||
return {
|
return {
|
||||||
@@ -40,3 +43,19 @@ test("AppRuntime also installs EffectLogger through Observability.layer", async
|
|||||||
expect(current.effectLogger).toBe(true)
|
expect(current.effectLogger).toBe(true)
|
||||||
expect(current.defaultLogger).toBe(false)
|
expect(current.defaultLogger).toBe(false)
|
||||||
})
|
})
|
||||||
|
|
||||||
|
test("AppRuntime attaches InstanceRef from ALS", async () => {
|
||||||
|
await using tmp = await tmpdir({ git: true })
|
||||||
|
|
||||||
|
const dir = await Instance.provide({
|
||||||
|
directory: tmp.path,
|
||||||
|
fn: () =>
|
||||||
|
AppRuntime.runPromise(
|
||||||
|
Effect.gen(function* () {
|
||||||
|
return (yield* InstanceRef)?.directory
|
||||||
|
}),
|
||||||
|
),
|
||||||
|
})
|
||||||
|
|
||||||
|
expect(dir).toBe(tmp.path)
|
||||||
|
})
|
||||||
|
|||||||
@@ -483,6 +483,48 @@ it.live("loop continues when finish is tool-calls", () =>
|
|||||||
),
|
),
|
||||||
)
|
)
|
||||||
|
|
||||||
|
it.live("glob tool keeps instance context during prompt runs", () =>
|
||||||
|
provideTmpdirServer(
|
||||||
|
({ dir, llm }) =>
|
||||||
|
Effect.gen(function* () {
|
||||||
|
const prompt = yield* SessionPrompt.Service
|
||||||
|
const sessions = yield* Session.Service
|
||||||
|
const session = yield* sessions.create({
|
||||||
|
title: "Glob context",
|
||||||
|
permission: [{ permission: "*", pattern: "*", action: "allow" }],
|
||||||
|
})
|
||||||
|
const file = path.join(dir, "probe.txt")
|
||||||
|
yield* Effect.promise(() => Bun.write(file, "probe"))
|
||||||
|
|
||||||
|
yield* prompt.prompt({
|
||||||
|
sessionID: session.id,
|
||||||
|
agent: "build",
|
||||||
|
noReply: true,
|
||||||
|
parts: [{ type: "text", text: "find text files" }],
|
||||||
|
})
|
||||||
|
yield* llm.tool("glob", { pattern: "**/*.txt" })
|
||||||
|
yield* llm.text("done")
|
||||||
|
|
||||||
|
const result = yield* prompt.loop({ sessionID: session.id })
|
||||||
|
expect(result.info.role).toBe("assistant")
|
||||||
|
|
||||||
|
const msgs = yield* MessageV2.filterCompactedEffect(session.id)
|
||||||
|
const tool = msgs
|
||||||
|
.flatMap((msg) => msg.parts)
|
||||||
|
.find(
|
||||||
|
(part): part is CompletedToolPart =>
|
||||||
|
part.type === "tool" && part.tool === "glob" && part.state.status === "completed",
|
||||||
|
)
|
||||||
|
if (!tool) return
|
||||||
|
|
||||||
|
expect(tool.state.output).toContain(file)
|
||||||
|
expect(tool.state.output).not.toContain("No context found for instance")
|
||||||
|
expect(result.parts.some((part) => part.type === "text" && part.text === "done")).toBe(true)
|
||||||
|
}),
|
||||||
|
{ git: true, config: providerCfg },
|
||||||
|
),
|
||||||
|
)
|
||||||
|
|
||||||
it.live("loop continues when finish is stop but assistant has tool parts", () =>
|
it.live("loop continues when finish is stop but assistant has tool parts", () =>
|
||||||
provideTmpdirServer(
|
provideTmpdirServer(
|
||||||
Effect.fnUntraced(function* ({ llm }) {
|
Effect.fnUntraced(function* ({ llm }) {
|
||||||
|
|||||||
Reference in New Issue
Block a user