feat(ultrawork): implement per-message model override with deferred DB retry strategy
- Add per-message ultrawork mode detection via keyword matching - Implement deferred DB override strategy using microtask retry loop - Fall back to setTimeout after 10 microtask retries for robustness - Update agent configuration schema with ultrawork model/variant fields - Integrate with chat.message hook to apply overrides on detection - Add comprehensive tests for all override scenarios - Generated schema includes ultrawork configuration 🤖 Generated with assistance of OhMyOpenCode (https://github.com/code-yeongyu/oh-my-opencode)
This commit is contained in:
YeonGyu-Kim
@@ -0,0 +1,180 @@
|
||||
import { describe, expect, test, beforeEach, afterEach, spyOn } from "bun:test"
|
||||
import { Database } from "bun:sqlite"
|
||||
import { mkdtempSync, mkdirSync, rmSync } from "node:fs"
|
||||
import { join } from "node:path"
|
||||
import { tmpdir } from "node:os"
|
||||
import * as dataPathModule from "../shared/data-path"
|
||||
import * as sharedModule from "../shared"
|
||||
|
||||
function flushMicrotasks(depth: number): Promise<void> {
|
||||
return new Promise<void>((resolve) => {
|
||||
let remaining = depth
|
||||
function step() {
|
||||
if (remaining <= 0) { resolve(); return }
|
||||
remaining--
|
||||
queueMicrotask(step)
|
||||
}
|
||||
queueMicrotask(step)
|
||||
})
|
||||
}
|
||||
|
||||
function flushWithTimeout(): Promise<void> {
|
||||
return new Promise<void>((resolve) => setTimeout(resolve, 10))
|
||||
}
|
||||
|
||||
describe("scheduleDeferredModelOverride", () => {
|
||||
let tempDir: string
|
||||
let dbPath: string
|
||||
let logSpy: ReturnType<typeof spyOn>
|
||||
let getDataDirSpy: ReturnType<typeof spyOn>
|
||||
|
||||
beforeEach(() => {
|
||||
tempDir = mkdtempSync(join(tmpdir(), "ultrawork-db-test-"))
|
||||
const opencodePath = join(tempDir, "opencode")
|
||||
mkdirSync(opencodePath, { recursive: true })
|
||||
dbPath = join(opencodePath, "opencode.db")
|
||||
|
||||
const db = new Database(dbPath)
|
||||
db.run(`
|
||||
CREATE TABLE IF NOT EXISTS message (
|
||||
id TEXT PRIMARY KEY,
|
||||
session_id TEXT NOT NULL,
|
||||
time_created TEXT NOT NULL DEFAULT (datetime('now')),
|
||||
time_updated TEXT NOT NULL DEFAULT (datetime('now')),
|
||||
data TEXT NOT NULL DEFAULT '{}'
|
||||
)
|
||||
`)
|
||||
db.close()
|
||||
|
||||
getDataDirSpy = spyOn(dataPathModule, "getDataDir").mockReturnValue(tempDir)
|
||||
logSpy = spyOn(sharedModule, "log").mockImplementation(() => {})
|
||||
})
|
||||
|
||||
afterEach(() => {
|
||||
getDataDirSpy?.mockRestore()
|
||||
logSpy?.mockRestore()
|
||||
rmSync(tempDir, { recursive: true, force: true })
|
||||
})
|
||||
|
||||
function insertMessage(id: string, model: { providerID: string; modelID: string }) {
|
||||
const db = new Database(dbPath)
|
||||
db.run(
|
||||
`INSERT INTO message (id, session_id, data) VALUES (?, ?, ?)`,
|
||||
id,
|
||||
"ses_test",
|
||||
JSON.stringify({ model }),
|
||||
)
|
||||
db.close()
|
||||
}
|
||||
|
||||
function readMessageModel(id: string): { providerID: string; modelID: string } | null {
|
||||
const db = new Database(dbPath)
|
||||
const row = db.query(`SELECT data FROM message WHERE id = ?`).get(id) as
|
||||
| { data: string }
|
||||
| null
|
||||
db.close()
|
||||
if (!row) return null
|
||||
const parsed = JSON.parse(row.data)
|
||||
return parsed.model ?? null
|
||||
}
|
||||
|
||||
function readMessageField(id: string, field: string): unknown {
|
||||
const db = new Database(dbPath)
|
||||
const row = db.query(`SELECT data FROM message WHERE id = ?`).get(id) as
|
||||
| { data: string }
|
||||
| null
|
||||
db.close()
|
||||
if (!row) return null
|
||||
return JSON.parse(row.data)[field] ?? null
|
||||
}
|
||||
|
||||
test("should update model in DB after microtask flushes", async () => {
|
||||
//#given
|
||||
insertMessage("msg_001", { providerID: "anthropic", modelID: "claude-sonnet-4-6" })
|
||||
|
||||
//#when
|
||||
const { scheduleDeferredModelOverride } = await import("./ultrawork-db-model-override")
|
||||
scheduleDeferredModelOverride(
|
||||
"msg_001",
|
||||
{ providerID: "anthropic", modelID: "claude-opus-4-6" },
|
||||
)
|
||||
await flushMicrotasks(5)
|
||||
|
||||
//#then
|
||||
const model = readMessageModel("msg_001")
|
||||
expect(model).toEqual({ providerID: "anthropic", modelID: "claude-opus-4-6" })
|
||||
})
|
||||
|
||||
test("should update variant and thinking fields when variant provided", async () => {
|
||||
//#given
|
||||
insertMessage("msg_002", { providerID: "anthropic", modelID: "claude-sonnet-4-6" })
|
||||
|
||||
//#when
|
||||
const { scheduleDeferredModelOverride } = await import("./ultrawork-db-model-override")
|
||||
scheduleDeferredModelOverride(
|
||||
"msg_002",
|
||||
{ providerID: "anthropic", modelID: "claude-opus-4-6" },
|
||||
"max",
|
||||
)
|
||||
await flushMicrotasks(5)
|
||||
|
||||
//#then
|
||||
expect(readMessageField("msg_002", "variant")).toBe("max")
|
||||
expect(readMessageField("msg_002", "thinking")).toBe("max")
|
||||
})
|
||||
|
||||
test("should fall back to setTimeout when message never appears", async () => {
|
||||
//#given — no message inserted
|
||||
|
||||
//#when
|
||||
const { scheduleDeferredModelOverride } = await import("./ultrawork-db-model-override")
|
||||
scheduleDeferredModelOverride(
|
||||
"msg_nonexistent",
|
||||
{ providerID: "anthropic", modelID: "claude-opus-4-6" },
|
||||
)
|
||||
await flushWithTimeout()
|
||||
|
||||
//#then
|
||||
expect(logSpy).toHaveBeenCalledWith(
|
||||
expect.stringContaining("setTimeout fallback failed"),
|
||||
expect.objectContaining({ messageId: "msg_nonexistent" }),
|
||||
)
|
||||
})
|
||||
|
||||
test("should not update variant fields when variant is undefined", async () => {
|
||||
//#given
|
||||
insertMessage("msg_003", { providerID: "anthropic", modelID: "claude-sonnet-4-6" })
|
||||
|
||||
//#when
|
||||
const { scheduleDeferredModelOverride } = await import("./ultrawork-db-model-override")
|
||||
scheduleDeferredModelOverride(
|
||||
"msg_003",
|
||||
{ providerID: "anthropic", modelID: "claude-opus-4-6" },
|
||||
)
|
||||
await flushMicrotasks(5)
|
||||
|
||||
//#then
|
||||
const model = readMessageModel("msg_003")
|
||||
expect(model).toEqual({ providerID: "anthropic", modelID: "claude-opus-4-6" })
|
||||
expect(readMessageField("msg_003", "variant")).toBeNull()
|
||||
expect(readMessageField("msg_003", "thinking")).toBeNull()
|
||||
})
|
||||
|
||||
test("should not crash when DB path does not exist", async () => {
|
||||
//#given
|
||||
getDataDirSpy.mockReturnValue("/nonexistent/path/that/does/not/exist")
|
||||
|
||||
//#when
|
||||
const { scheduleDeferredModelOverride } = await import("./ultrawork-db-model-override")
|
||||
scheduleDeferredModelOverride(
|
||||
"msg_004",
|
||||
{ providerID: "anthropic", modelID: "claude-opus-4-6" },
|
||||
)
|
||||
await flushMicrotasks(5)
|
||||
|
||||
//#then
|
||||
expect(logSpy).toHaveBeenCalledWith(
|
||||
expect.stringContaining("DB not found"),
|
||||
)
|
||||
})
|
||||
})
|
||||
Reference in New Issue
Block a user