refactor: simplify instance context helpers in prompt tests

pull/20304/head
Kit Langton 2026-03-31 16:32:12 -04:00
parent 4ff0fbc043
commit 90469bbb7e
7 changed files with 110 additions and 172 deletions

View File

@ -1,14 +0,0 @@
import { Fiber } from "effect"
import * as ServiceMap from "effect/ServiceMap"
import { Instance } from "@/project/instance"
import { InstanceRef } from "./instance-ref"
export function bind<F extends (...args: any[]) => any>(fn: F): F {
try {
return Instance.bind(fn)
} catch {}
const fiber = Fiber.getCurrent()
const ctx = fiber ? ServiceMap.getReferenceUnsafe(fiber.services, InstanceRef) : undefined
if (!ctx) return fn
return ((...args: any[]) => Instance.restore(ctx, () => fn(...args))) as F
}

View File

@ -1,6 +1,5 @@
import { Effect, ScopedCache, Scope } from "effect"
import { Effect, Fiber, ScopedCache, Scope, ServiceMap } from "effect"
import { Instance, type InstanceContext } from "@/project/instance"
import { bind as bindInstance } from "./instance-bind"
import { InstanceRef } from "./instance-ref"
import { registerDisposer } from "./instance-registry"
@ -12,17 +11,24 @@ export interface InstanceState<A, E = never, R = never> {
}
export namespace InstanceState {
export const bind = bindInstance
export const bind = <F extends (...args: any[]) => any>(fn: F): F => {
try {
return Instance.bind(fn)
} catch {}
const fiber = Fiber.getCurrent()
const ctx = fiber ? ServiceMap.getReferenceUnsafe(fiber.services, InstanceRef) : undefined
if (!ctx) return fn
return ((...args: any[]) => Instance.restore(ctx, () => fn(...args))) as F
}
export const context = Effect.gen(function* () {
const ref = yield* InstanceRef
return ref ?? Instance.current
})
export const context = Effect.fnUntraced(function* () {
return (yield* InstanceRef) ?? Instance.current
})()
export const directory = Effect.gen(function* () {
const ref = yield* InstanceRef
return ref ? ref.directory : Instance.directory
})
export const directory = Effect.fnUntraced(function* () {
const ctx = yield* InstanceRef
return ctx ? ctx.directory : Instance.directory
})()
export const make = <A, E = never, R = never>(
init: (ctx: InstanceContext) => Effect.Effect<A, E, R | Scope.Scope>,
@ -31,9 +37,9 @@ export namespace InstanceState {
const cache = yield* ScopedCache.make<string, A, E, R>({
capacity: Number.POSITIVE_INFINITY,
lookup: () =>
Effect.gen(function* () {
Effect.fnUntraced(function* () {
return yield* init(yield* context)
}),
})(),
})
const off = registerDisposer((directory) => Effect.runPromise(ScopedCache.invalidate(cache, directory)))

View File

@ -5,7 +5,7 @@ import { InstanceRef } from "./instance-ref"
export const memoMap = Layer.makeMemoMapUnsafe()
function provide<A, E, R>(effect: Effect.Effect<A, E, R>): Effect.Effect<A, E, R> {
function attach<A, E, R>(effect: Effect.Effect<A, E, R>): Effect.Effect<A, E, R> {
try {
const ctx = Instance.current
return Effect.provideService(effect, InstanceRef, ctx)
@ -18,13 +18,13 @@ export function makeRuntime<I, S, E>(service: ServiceMap.Service<I, S>, layer: L
const getRuntime = () => (rt ??= ManagedRuntime.make(layer, { memoMap }))
return {
runSync: <A, Err>(fn: (svc: S) => Effect.Effect<A, Err, I>) => getRuntime().runSync(provide(service.use(fn))),
runSync: <A, Err>(fn: (svc: S) => Effect.Effect<A, Err, I>) => getRuntime().runSync(attach(service.use(fn))),
runPromiseExit: <A, Err>(fn: (svc: S) => Effect.Effect<A, Err, I>, options?: Effect.RunOptions) =>
getRuntime().runPromiseExit(provide(service.use(fn)), options),
getRuntime().runPromiseExit(attach(service.use(fn)), options),
runPromise: <A, Err>(fn: (svc: S) => Effect.Effect<A, Err, I>, options?: Effect.RunOptions) =>
getRuntime().runPromise(provide(service.use(fn)), options),
runFork: <A, Err>(fn: (svc: S) => Effect.Effect<A, Err, I>) => getRuntime().runFork(provide(service.use(fn))),
getRuntime().runPromise(attach(service.use(fn)), options),
runFork: <A, Err>(fn: (svc: S) => Effect.Effect<A, Err, I>) => getRuntime().runFork(attach(service.use(fn))),
runCallback: <A, Err>(fn: (svc: S) => Effect.Effect<A, Err, I>) =>
getRuntime().runCallback(provide(service.use(fn))),
getRuntime().runCallback(attach(service.use(fn))),
}
}

View File

@ -498,10 +498,10 @@ export namespace Session {
permission?: Permission.Ruleset
workspaceID?: WorkspaceID
}) {
const dir = yield* InstanceState.directory
const directory = yield* InstanceState.directory
return yield* createNext({
parentID: input?.parentID,
directory: dir,
directory,
title: input?.title,
permission: input?.permission,
workspaceID: input?.workspaceID,
@ -509,11 +509,11 @@ export namespace Session {
})
const fork = Effect.fn("Session.fork")(function* (input: { sessionID: SessionID; messageID?: MessageID }) {
const dir = yield* InstanceState.directory
const directory = yield* InstanceState.directory
const original = yield* get(input.sessionID)
const title = getForkedTitle(original.title)
const session = yield* createNext({
directory: dir,
directory,
workspaceID: original.workspaceID,
title,
})

View File

@ -12,7 +12,7 @@ import path from "path"
import { readFileSync, readdirSync, existsSync } from "fs"
import { Flag } from "../flag/flag"
import { CHANNEL } from "../installation/meta"
import { bind } from "@/effect/instance-bind"
import { InstanceState } from "@/effect/instance-state"
import { iife } from "@/util/iife"
import { init } from "#db"
@ -142,7 +142,7 @@ export namespace Database {
}
export function effect(fn: () => any | Promise<any>) {
const bound = bind(fn)
const bound = InstanceState.bind(fn)
try {
ctx.use().effects.push(bound)
} catch {
@ -163,7 +163,7 @@ export namespace Database {
} catch (err) {
if (err instanceof Context.NotFound) {
const effects: (() => void | Promise<void>)[] = []
const txCallback = bind((tx: TxOrDb) => ctx.provide({ tx, effects }, () => callback(tx)))
const txCallback = InstanceState.bind((tx: TxOrDb) => ctx.provide({ tx, effects }, () => callback(tx)))
const result = Client().transaction(txCallback, { behavior: options?.behavior })
for (const effect of effects) effect()
return result as NotPromise<T>

View File

@ -340,6 +340,84 @@ it.live("loop calls LLM and returns assistant message", () =>
),
)
it.live("static loop returns assistant text through local provider", () =>
provideTmpdirServer(
Effect.fnUntraced(function* ({ llm }) {
const session = yield* Effect.promise(() =>
Session.create({
title: "Prompt provider",
permission: [{ permission: "*", pattern: "*", action: "allow" }],
}),
)
yield* Effect.promise(() =>
SessionPrompt.prompt({
sessionID: session.id,
agent: "build",
noReply: true,
parts: [{ type: "text", text: "hello" }],
}),
)
yield* llm.text("world")
const result = yield* Effect.promise(() => SessionPrompt.loop({ sessionID: session.id }))
expect(result.info.role).toBe("assistant")
expect(result.parts.some((part) => part.type === "text" && part.text === "world")).toBe(true)
expect(yield* llm.hits).toHaveLength(1)
expect(yield* llm.pending).toBe(0)
}),
{ git: true, config: providerCfg },
),
)
it.live("static loop consumes queued replies across turns", () =>
provideTmpdirServer(
Effect.fnUntraced(function* ({ llm }) {
const session = yield* Effect.promise(() =>
Session.create({
title: "Prompt provider turns",
permission: [{ permission: "*", pattern: "*", action: "allow" }],
}),
)
yield* Effect.promise(() =>
SessionPrompt.prompt({
sessionID: session.id,
agent: "build",
noReply: true,
parts: [{ type: "text", text: "hello one" }],
}),
)
yield* llm.text("world one")
const first = yield* Effect.promise(() => SessionPrompt.loop({ sessionID: session.id }))
expect(first.info.role).toBe("assistant")
expect(first.parts.some((part) => part.type === "text" && part.text === "world one")).toBe(true)
yield* Effect.promise(() =>
SessionPrompt.prompt({
sessionID: session.id,
agent: "build",
noReply: true,
parts: [{ type: "text", text: "hello two" }],
}),
)
yield* llm.text("world two")
const second = yield* Effect.promise(() => SessionPrompt.loop({ sessionID: session.id }))
expect(second.info.role).toBe("assistant")
expect(second.parts.some((part) => part.type === "text" && part.text === "world two")).toBe(true)
expect(yield* llm.hits).toHaveLength(2)
expect(yield* llm.pending).toBe(0)
}),
{ git: true, config: providerCfg },
),
)
it.live("loop continues when finish is tool-calls", () =>
provideTmpdirServer(
Effect.fnUntraced(function* ({ llm }) {

View File

@ -1,132 +0,0 @@
import { describe, expect } from "bun:test"
import { Effect } from "effect"
import { NodeFileSystem } from "@effect/platform-node"
import * as CrossSpawnSpawner from "../../src/effect/cross-spawn-spawner"
import { Session } from "../../src/session"
import { SessionPrompt } from "../../src/session/prompt"
import { Log } from "../../src/util/log"
import { testEffect } from "../lib/effect"
import { provideTmpdirServer } from "../fixture/fixture"
import { TestLLMServer } from "../lib/llm-server"
import { Layer } from "effect"
Log.init({ print: false })
const baseLayer = Layer.mergeAll(NodeFileSystem.layer, CrossSpawnSpawner.defaultLayer, TestLLMServer.layer)
const it = testEffect(baseLayer)
function makeConfig(url: string) {
return {
provider: {
test: {
name: "Test",
env: [],
npm: "@ai-sdk/openai-compatible",
models: {
"gpt-5-nano": {
id: "gpt-5-nano",
name: "Test Model",
attachment: false,
reasoning: false,
temperature: false,
tool_call: true,
release_date: "2025-01-01",
limit: { context: 100000, output: 10000 },
cost: { input: 0, output: 0 },
options: {},
},
},
options: {
apiKey: "test-key",
baseURL: url,
},
},
},
agent: {
build: {
model: "test/gpt-5-nano",
},
},
}
}
describe("session.prompt provider integration", () => {
it.live("loop returns assistant text through local provider", () =>
provideTmpdirServer(
Effect.fnUntraced(function* ({ llm }) {
const session = yield* Effect.promise(() =>
Session.create({
title: "Prompt provider",
permission: [{ permission: "*", pattern: "*", action: "allow" }],
}),
)
yield* Effect.promise(() =>
SessionPrompt.prompt({
sessionID: session.id,
agent: "build",
noReply: true,
parts: [{ type: "text", text: "hello" }],
}),
)
yield* llm.text("world")
const result = yield* Effect.promise(() => SessionPrompt.loop({ sessionID: session.id }))
expect(result.info.role).toBe("assistant")
expect(result.parts.some((part) => part.type === "text" && part.text === "world")).toBe(true)
expect(yield* llm.hits).toHaveLength(1)
expect(yield* llm.pending).toBe(0)
}),
{ git: true, config: makeConfig },
),
)
it.live("loop consumes queued replies across turns", () =>
provideTmpdirServer(
Effect.fnUntraced(function* ({ llm }) {
const session = yield* Effect.promise(() =>
Session.create({
title: "Prompt provider turns",
permission: [{ permission: "*", pattern: "*", action: "allow" }],
}),
)
yield* Effect.promise(() =>
SessionPrompt.prompt({
sessionID: session.id,
agent: "build",
noReply: true,
parts: [{ type: "text", text: "hello one" }],
}),
)
yield* llm.text("world one")
const first = yield* Effect.promise(() => SessionPrompt.loop({ sessionID: session.id }))
expect(first.info.role).toBe("assistant")
expect(first.parts.some((part) => part.type === "text" && part.text === "world one")).toBe(true)
yield* Effect.promise(() =>
SessionPrompt.prompt({
sessionID: session.id,
agent: "build",
noReply: true,
parts: [{ type: "text", text: "hello two" }],
}),
)
yield* llm.text("world two")
const second = yield* Effect.promise(() => SessionPrompt.loop({ sessionID: session.id }))
expect(second.info.role).toBe("assistant")
expect(second.parts.some((part) => part.type === "text" && part.text === "world two")).toBe(true)
expect(yield* llm.hits).toHaveLength(2)
expect(yield* llm.pending).toBe(0)
}),
{ git: true, config: makeConfig },
),
)
})