@batonfx/foldkit 0.1.0

package/src/connection.ts

@@ -0,0 +1,133 @@
+import { Cause, Context, Effect, Layer, Option, Ref, Result, Schema, Stream } from "effect"
+import { Socket } from "effect/unstable/socket"
+import { m } from "foldkit/message"
+import type { CallableTaggedStruct } from "foldkit/schema"
+import { Client, Wire } from "@batonfx/transport"
+
+/** @experimental */
+export const ConnectionOpened: CallableTaggedStruct<"ConnectionOpened", {}> = m("ConnectionOpened")
+
+/** @experimental */
+export const ConnectionLost: CallableTaggedStruct<"ConnectionLost", {}> = m("ConnectionLost")
+
+/** @experimental */
+export const ConnectionFailed: CallableTaggedStruct<"ConnectionFailed", { reason: typeof Schema.String }> = m(
+  "ConnectionFailed",
+  { reason: Schema.String },
+)
+
+/** @experimental */
+export type Incoming =
+  | Wire.LooseServerFrameType
+  | typeof ConnectionOpened.Type
+  | typeof ConnectionLost.Type
+  | typeof ConnectionFailed.Type
+
+/** @experimental */
+export const Incoming: Schema.Schema<Incoming> = Schema.Union([
+  Wire.LooseServerFrame,
+  ConnectionOpened,
+  ConnectionLost,
+  ConnectionFailed,
+])
+
+/** @experimental */
+export class SendFailed extends Schema.TaggedErrorClass<SendFailed>()("@batonfx/foldkit/SendFailed", {
+  reason: Schema.String,
+}) {}
+
+/** @experimental */
+export interface Interface {
+  readonly frames: (options: {
+    readonly sessionId: string
+    readonly afterSeq?: number
+  }) => Stream.Stream<Incoming, never>
+  readonly send: (frame: Wire.ClientFrameType) => Effect.Effect<void, SendFailed>
+}
+
+/** @experimental */
+export class AgentConnection extends Context.Service<AgentConnection, Interface>()(
+  "@batonfx/foldkit/AgentConnection",
+) {}
+
+interface ActiveConnection {
+  readonly sessionId: string
+  readonly connection: Client.Connection
+}
+
+const reasonFrom = (error: unknown): string => {
+  if (error instanceof SendFailed) return error.reason
+  if (error instanceof Error) return error.message
+  return String(error)
+}
+
+const statusIncoming = (status: Client.ConnectionStatus): Option.Option<Incoming> => {
+  switch (status._tag) {
+    case "Open":
+      return Option.some(ConnectionOpened())
+    case "Reconnecting":
+    case "Closed":
+      return Option.some(ConnectionLost())
+    case "Connecting":
+      return Option.none()
+  }
+}
+
+/** @experimental */
+export const testLayer = (implementation: Interface): Layer.Layer<AgentConnection> =>
+  Layer.succeed(AgentConnection, AgentConnection.of(implementation))
+
+/** @experimental */
+export const layerWebSocket = (options: {
+  readonly url: string
+}): Layer.Layer<AgentConnection, never, Socket.WebSocketConstructor> =>
+  Layer.effect(
+    AgentConnection,
+    Effect.gen(function* () {
+      const client = yield* Client.AgentClient
+      const active = yield* Ref.make<Option.Option<ActiveConnection>>(Option.none())
+
+      const clearActive = (connection: Client.Connection) =>
+        Ref.update(active, (current) =>
+          Option.isSome(current) && current.value.connection === connection ? Option.none() : current,
+        )
+
+      return AgentConnection.of({
+        frames: ({ sessionId }) =>
+          Stream.unwrap(
+            Effect.gen(function* () {
+              const connection = yield* client.connect({ url: options.url, sessionId })
+              yield* Ref.set(active, Option.some({ sessionId, connection }))
+              yield* Effect.addFinalizer(() => clearActive(connection))
+              const statuses = connection.status.pipe(
+                Stream.filterMap((status) =>
+                  Option.match(statusIncoming(status), {
+                    onNone: () => Result.fail(undefined),
+                    onSome: Result.succeed,
+                  }),
+                ),
+              )
+              const frames = connection.frames.pipe(
+                Stream.map((frame): Incoming => frame),
+                Stream.catchCause((cause) =>
+                  Stream.succeed(ConnectionFailed({ reason: reasonFrom(Cause.squash(cause)) })),
+                ),
+              )
+              return statuses.pipe(Stream.merge(frames))
+            }),
+          ),
+        send: (frame) =>
+          Ref.get(active).pipe(
+            Effect.flatMap(
+              Option.match({
+                onNone: () => Effect.fail(new SendFailed({ reason: "No active agent connection" })),
+                onSome: ({ connection }) =>
+                  connection
+                    .send(frame)
+                    .pipe(Effect.mapError((error) => new SendFailed({ reason: reasonFrom(error) }))),
+              }),
+            ),
+          ),
+      })
+    }),
+  ).pipe(Layer.provide(Client.layerWebSocket))

package/src/index.ts

@@ -0,0 +1,2 @@
+export * as Chat from "./chat"
+export * as Connection from "./connection"

package/test/chat.test.ts

@@ -0,0 +1,146 @@
+import { describe, expect, it } from "@effect/vitest"
+import { Effect, Option, Schema, Stream } from "effect"
+import * as Ai from "effect/unstable/ai"
+import { Wire } from "@batonfx/transport"
+import { Chat, Connection } from "../src/index"
+
+const eventFrame = (seq: number, event: Wire.EventType): Wire.LooseServerFrameType => ({ _tag: "Event", seq, event })
+
+const updateWith = (model: Chat.Model, incoming: Connection.Incoming) =>
+  Chat.update(model, Chat.ReceivedAgent({ incoming }))
+
+describe("Chat", () => {
+  it("folds a full run into display entries and a completion out-message", () => {
+    let model = Chat.initialModel("s-chat")
+    let out: Option.Option<Chat.OutMessage> = Option.none()
+
+    ;[model, , out] = updateWith(model, eventFrame(0, { _tag: "TurnStarted", turn: 0 }))
+    expect(model.run).toEqual({ _tag: "Running", turn: 0 })
+    expect(model.streaming).toEqual({ turn: 0, text: "", reasoning: "" })
+    expect(Option.isNone(out)).toBe(true)
+    ;[model] = updateWith(
+      model,
+      eventFrame(1, {
+        _tag: "ModelPart",
+        turn: 0,
+        part: Ai.Response.makePart("text-delta", { id: "text-1", delta: "Hello " }),
+      }),
+    )
+    ;[model] = updateWith(
+      model,
+      eventFrame(2, {
+        _tag: "ModelPart",
+        turn: 0,
+        part: Ai.Response.makePart("text-delta", { id: "text-1", delta: "world" }),
+      }),
+    )
+
+    const call = Ai.Response.makePart("tool-call", {
+      id: "call-1",
+      name: "lookup",
+      params: { q: "baton" },
+      providerExecuted: false,
+    })
+    ;[model] = updateWith(model, eventFrame(3, { _tag: "ModelPart", turn: 0, part: call }))
+    expect(model.entries).toEqual([
+      {
+        _tag: "ToolEntry",
+        callId: "call-1",
+        name: "lookup",
+        params: { q: "baton" },
+        outcome: { _tag: "Pending" },
+        progress: [],
+      },
+    ])
+
+    const result = Ai.Response.makePart("tool-result", {
+      id: "call-1",
+      name: "lookup",
+      result: { ok: true },
+      encodedResult: { ok: true },
+      isFailure: false,
+      providerExecuted: false,
+      preliminary: false,
+    })
+    ;[model] = updateWith(model, eventFrame(4, { _tag: "ToolExecutionCompleted", turn: 0, call, result }))
+    expect(model.entries[0]).toMatchObject({
+      _tag: "ToolEntry",
+      outcome: { _tag: "Completed", isFailure: false, result: { ok: true } },
+    })
+    ;[model] = updateWith(model, eventFrame(5, { _tag: "TurnCompleted", turn: 0 }))
+    expect(model.streaming).toBeNull()
+    expect(model.entries[1]).toEqual({ _tag: "AssistantEntry", text: "Hello world", reasoning: null })
+    ;[model, , out] = updateWith(model, eventFrame(6, { _tag: "Completed", turns: 1, text: "Done" }))
+    expect(model.run).toEqual({ _tag: "Idle" })
+    expect(Option.getOrUndefined(out)).toEqual({ _tag: "RunCompleted", text: "Done" })
+  })
+
+  it("drops replayed frames whose seq is not newer than lastSeq", () => {
+    const model = { ...Chat.initialModel("s-chat"), lastSeq: 5, entries: [Chat.UserEntry({ text: "already" })] }
+    const [next, commands, out] = updateWith(model, eventFrame(5, { _tag: "TurnStarted", turn: 1 }))
+
+    expect(next).toEqual(model)
+    expect(commands).toEqual([])
+    expect(Option.isNone(out)).toBe(true)
+  })
+
+  it("surfaces approval suspension and emits approval commands", () => {
+    let model = Chat.initialModel("s-chat")
+    let out: Option.Option<Chat.OutMessage> = Option.none()
+    ;[model, , out] = updateWith(
+      model,
+      Schema.decodeUnknownSync(Wire.LooseServerFrame)({
+        _tag: "Suspended",
+        seq: 0,
+        suspension: {
+          _tag: "@batonfx/core/AgentSuspended",
+          token: "approval-token",
+          reason: "approval",
+          tool_call_id: "call-approval",
+          tool_name: "lookup",
+          tool_params: { q: "baton" },
+        },
+      }),
+    )
+
+    expect(model.run).toEqual({
+      _tag: "AwaitingApproval",
+      token: "approval-token",
+      toolName: "lookup",
+      params: { q: "baton" },
+    })
+    expect(Option.getOrUndefined(out)).toEqual({ _tag: "ApprovalRequired" })
+
+    const [, commands] = Chat.update(model, Chat.ClickedApprove())
+    expect(commands).toHaveLength(1)
+    expect(commands[0]?.name).toBe("ResolveApproval")
+    expect(commands[0]?.args).toEqual({
+      sessionId: "s-chat",
+      token: "approval-token",
+      approved: true,
+      reason: null,
+    })
+  })
+
+  it.effect("converts command send failures into FailedAgentCommand messages", () =>
+    Effect.gen(function* () {
+      let model = Chat.initialModel("s-chat")
+      ;[model] = Chat.update(model, Chat.ChangedDraft({ text: "hello" }))
+      const [next, commands] = Chat.update(model, Chat.SubmittedMessage())
+
+      expect(next.draft).toBe("")
+      expect(next.entries).toEqual([Chat.UserEntry({ text: "hello" })])
+      expect(commands).toHaveLength(1)
+      const message = yield* commands[0]!.effect
+
+      expect(message).toEqual({ _tag: "FailedAgentCommand", reason: "offline" })
+    }).pipe(
+      Effect.provide(
+        Connection.testLayer({
+          frames: () => Stream.empty,
+          send: () => Effect.fail(new Connection.SendFailed({ reason: "offline" })),
+        }),
+      ),
+    ),
+  )
+})

package/test/connection.test.ts

@@ -0,0 +1,35 @@
+import { describe, expect, it } from "@effect/vitest"
+import { Effect, Stream } from "effect"
+import { Chat, Connection } from "../src/index"
+
+describe("Connection", () => {
+  it.effect("testLayer provides AgentConnection frames and send", () =>
+    Effect.gen(function* () {
+      const incoming = Connection.ConnectionOpened()
+      const frames = yield* Connection.AgentConnection.use((connection) =>
+        connection.frames({ sessionId: "s", afterSeq: 1 }).pipe(Stream.runCollect),
+      )
+      yield* Connection.AgentConnection.use((connection) => connection.send({ _tag: "Cancel", sessionId: "s" }))
+
+      expect(frames).toEqual([incoming])
+    }).pipe(
+      Effect.provide(
+        Connection.testLayer({
+          frames: () => Stream.fromIterable([Connection.ConnectionOpened()]),
+          send: () => Effect.void,
+        }),
+      ),
+    ),
+  )
+
+  it("keeps the chat subscription alive across afterSeq-only changes", () => {
+    const subscription = Chat.subscriptions.agentFrames
+
+    expect(subscription.keepAliveEquivalence?.({ sessionId: "s", afterSeq: 1 }, { sessionId: "s", afterSeq: 2 })).toBe(
+      true,
+    )
+    expect(
+      subscription.keepAliveEquivalence?.({ sessionId: "s-1", afterSeq: 2 }, { sessionId: "s-2", afterSeq: 2 }),
+    ).toBe(false)
+  })
+})

package/tsconfig.json

@@ -0,0 +1,7 @@
+{
+  "extends": "../../tsconfig.json",
+  "compilerOptions": {
+    "rootDir": "."
+  },
+  "include": ["src/**/*.ts", "test/**/*.ts"]
+}