diff --git a/apps/server/src/ws.test.ts b/apps/server/src/ws.test.ts index a1dfa5f4b868c..ac60becc2d78a 100644 --- a/apps/server/src/ws.test.ts +++ b/apps/server/src/ws.test.ts @@ -303,6 +303,74 @@ it.live.each(["thread", "shell"] as const)( }).pipe(Effect.scoped, Effect.provide(ReplayTestLayer)), ); +it.effect("coalesces streaming replay before synchronization and the live tail", () => + Effect.gen(function* () { + const store = yield* OrchestrationEventStore.OrchestrationEventStore; + const sql = yield* SqlClient.SqlClient; + const threadId = ThreadId.make("thread:streaming-replay"); + const now = "2026-10-27T00:00:00.000Z"; + const insert = (version: number, text: string, streaming: boolean) => sql<{ + sequence: number; + }>` + INSERT INTO orchestration_events ( + event_id, aggregate_kind, stream_id, stream_version, event_type, + occurred_at, actor_kind, payload_json, metadata_json, application_event_version + ) VALUES ( + ${`streaming-replay:${version}`}, 'thread', ${threadId}, ${version}, 'message.updated', + ${now}, 'server', ${JSON.stringify({ + createdBy: "agent", + creationSource: "web", + id: "message:streaming-replay", + threadId, + runId: null, + nodeId: null, + role: "assistant", + text, + attachments: [], + streaming, + createdAt: now, + updatedAt: now, + })}, '{}', 2 + ) RETURNING sequence + `; + yield* insert(0, "H", true); + yield* insert(1, "He", true); + const [latest] = yield* insert(2, "Hello", true); + const [completed] = yield* insert(3, "Hello!", false); + const items = yield* subscribeOrchestrationV2Thread({ + threadId, + afterSequence: 0, + requestCompletionMarker: true, + }).pipe( + Effect.flatMap(Stream.runCollect), + Effect.provideService(ThreadManagement.ThreadManagementService, { + ensureLegacyTranscript: () => Effect.void, + streamStoredEventsFrom: () => + Stream.unwrap( + insert(4, "Next", false).pipe( + Effect.map(() => + store.readAgentEvents({ threadId, afterSequence: completed!.sequence }), + ), + ), + ), + } as unknown as ThreadManagement.ThreadManagementService["Service"]), + ); + assert.deepEqual( + items.map((item) => + item.kind === "event" + ? [item.sequence, item.event.type === "message.updated" ? item.event.payload.text : null] + : [item.kind], + ), + [ + [latest!.sequence, "Hello"], + [completed!.sequence, "Hello!"], + ["synchronized"], + [completed!.sequence + 1, "Next"], + ], + ); + }).pipe(Effect.scoped, Effect.provide(ReplayTestLayer)), +); + it.effect("shell snapshots omit idle children and older settled threads only on request", () => Effect.gen(function* () { const shell = ( diff --git a/apps/server/src/ws.ts b/apps/server/src/ws.ts index 8d3666c2f9df0..e35aaf1bb93cd 100644 --- a/apps/server/src/ws.ts +++ b/apps/server/src/ws.ts @@ -142,7 +142,10 @@ import { type ShellApplicationEvent, } from "./orchestration-v2/ShellStream.ts"; import { ORCHESTRATION_V2_PROJECTION_SCHEMA_VERSION } from "./orchestration-v2/ProjectionStore.ts"; -import { coalesceThreadLiveStream } from "./orchestration-v2/ThreadLiveEventCoalescer.ts"; +import { + coalesceLiveUpdatedEvents, + coalesceThreadLiveStream, +} from "./orchestration-v2/ThreadLiveEventCoalescer.ts"; import { buildBoundedThreadStreamSnapshot, decideThreadResume, @@ -947,7 +950,9 @@ export const subscribeOrchestrationV2Thread = Effect.fn("ws.orchestrationV2.subs // Retained creation events may describe deleted threads; replace replay only if a snapshot exists. if (shell !== null) return yield* snapshotThenLive(); } - const replay = yield* loadReplayThrough(input.afterSequence, highWater); + const replay = coalesceLiveUpdatedEvents( + yield* loadReplayThrough(input.afterSequence, highWater), + ); const plan = decideThreadResume({ afterSequence: input.afterSequence, highWater, diff --git a/packages/client-runtime/src/authorization/layer.test.ts b/packages/client-runtime/src/authorization/layer.test.ts index d46878dedcd83..f7d8665549e7f 100644 --- a/packages/client-runtime/src/authorization/layer.test.ts +++ b/packages/client-runtime/src/authorization/layer.test.ts @@ -1,4 +1,8 @@ -import { AuthStandardClientScopes, EnvironmentId } from "@t3tools/contracts"; +import { + AuthStandardClientScopes, + EnvironmentId, + ORCHESTRATION_PROTOCOL_VERSION, +} from "@t3tools/contracts"; import { type RelayEnvironmentConnectResponse } from "@t3tools/contracts/relay"; import { describe, expect, it } from "@effect/vitest"; import * as Deferred from "effect/Deferred"; @@ -15,6 +19,11 @@ import * as RpcHttp from "../rpc/http.ts"; import * as ClientCapabilities from "../platform/capabilities.ts"; import * as RemoteEnvironmentAuthorization from "./service.ts"; import * as TokenStore from "./tokenStore.ts"; +import { BearerConnectionCredential, BearerConnectionProfile } from "../connection/catalog.ts"; +import * as ConnectionCredentialStore from "../connection/credentialStore.ts"; +import { BearerConnectionTarget } from "../connection/model.ts"; +import * as ConnectionProfileStore from "../connection/profileStore.ts"; +import * as ConnectionResolver from "../connection/resolver.ts"; const ENVIRONMENT_ID = EnvironmentId.make("environment-1"); const ENDPOINT = { @@ -186,7 +195,7 @@ const makeHarness = Effect.fn("TestRemoteAuthorization.makeHarness")(function* ( resetTokenCache: Effect.void, }); const layer = RemoteEnvironmentAuthorization.layer.pipe( - Layer.provide( + Layer.provideMerge( Layer.mergeAll( RpcHttp.layerRemoteHttpClient(fetch.fetchFn), Layer.succeed(ManagedRelay.ManagedRelayDpopSigner, signer), @@ -225,13 +234,124 @@ const makeHarness = Effect.fn("TestRemoteAuthorization.makeHarness")(function* ( }; }); +const prepareDirectBearer = Effect.fn("TestRemoteAuthorization.prepareDirectBearer")(function* ( + harness: Effect.Success>, + times: number, +) { + const target = new BearerConnectionTarget({ + environmentId: ENVIRONMENT_ID, + label: DESCRIPTOR.label, + connectionId: "direct-1", + }); + const entry = { + target, + profile: Option.some( + new BearerConnectionProfile({ + connectionId: target.connectionId, + environmentId: ENVIRONMENT_ID, + label: target.label, + httpBaseUrl: ENDPOINT.httpBaseUrl, + wsBaseUrl: ENDPOINT.wsBaseUrl, + }), + ), + enabled: true, + }; + const unused = () => Effect.die("unused"); + const layer = ConnectionResolver.layer.pipe( + Layer.provide( + Layer.mergeAll( + harness.layer, + Layer.succeed(ConnectionProfileStore.ConnectionProfileStore, { + get: unused, + put: unused, + remove: unused, + }), + Layer.succeed(ConnectionCredentialStore.ConnectionCredentialStore, { + get: () => Effect.succeedSome(new BearerConnectionCredential({ token: "bearer-token" })), + put: unused, + remove: unused, + }), + Layer.succeed(ClientCapabilities.PrimaryEnvironmentAuth, { + bearerToken: Effect.succeedNone, + }), + Layer.succeed(ClientCapabilities.SshEnvironmentGateway, { + provision: unused, + prepare: unused, + disconnect: unused, + }), + ), + ), + ); + const resolver = yield* ConnectionResolver.ConnectionResolver.pipe(Effect.provide(layer)); + return yield* Effect.forEach(Array.from({ length: times }), () => + resolver.prepare(entry).pipe(Effect.result), + ); +}); + describe("RemoteEnvironmentAuthorization", () => { - it.effect("reuses a validated bearer descriptor while issuing fresh websocket tickets", () => + it.effect("prepares a cold direct bearer connection with one descriptor request", () => + Effect.gen(function* () { + const harness = yield* makeHarness({ + responses: [ + Response.json({ + ...DESCRIPTOR, + orchestrationProtocolVersion: ORCHESTRATION_PROTOCOL_VERSION, + }), + websocketTicket("cold-ticket"), + ], + }); + + const [prepared] = yield* prepareDirectBearer(harness, 1); + + expect(prepared).toMatchObject({ + _tag: "Success", + success: { socketUrl: expect.stringContaining("wsTicket=cold-ticket") }, + }); + expect(harness.fetch.calls.map(([url]) => new URL(String(url)).pathname)).toEqual([ + "/.well-known/t3/environment", + "/api/auth/websocket-ticket/rpc", + ]); + }), + ); + + it.effect("blocks a warm direct bearer prepare when the environment changes identity", () => + Effect.gen(function* () { + const harness = yield* makeHarness({ + responses: [ + Response.json({ + ...DESCRIPTOR, + orchestrationProtocolVersion: ORCHESTRATION_PROTOCOL_VERSION, + }), + websocketTicket("first-ticket"), + Response.json({ + ...DESCRIPTOR, + environmentId: "another-environment", + orchestrationProtocolVersion: ORCHESTRATION_PROTOCOL_VERSION, + }), + ], + }); + + const [first, second] = yield* prepareDirectBearer(harness, 2); + + expect(first?._tag).toBe("Success"); + expect(second).toMatchObject({ + _tag: "Failure", + failure: { _tag: "ConnectionBlockedError", reason: "configuration" }, + }); + expect( + harness.fetch.calls.filter(([url]) => + String(url).endsWith("/api/auth/websocket-ticket/rpc"), + ), + ).toHaveLength(1); + }), + ); + it.effect("validates a fresh bearer descriptor while issuing fresh websocket tickets", () => Effect.gen(function* () { const harness = yield* makeHarness({ responses: [ Response.json(DESCRIPTOR), websocketTicket("first-ticket"), + Response.json(DESCRIPTOR), websocketTicket("second-ticket"), ], }); @@ -253,7 +373,7 @@ describe("RemoteEnvironmentAuthorization", () => { expect(second.socketUrl).toContain("wsTicket=second-ticket"); expect( harness.fetch.calls.filter(([url]) => String(url).endsWith("/.well-known/t3/environment")), - ).toHaveLength(1); + ).toHaveLength(2); expect( harness.fetch.calls.filter(([url]) => String(url).endsWith("/api/auth/websocket-ticket/rpc"), diff --git a/packages/client-runtime/src/authorization/service.ts b/packages/client-runtime/src/authorization/service.ts index d4c5ff9369616..2962393d7b4a2 100644 --- a/packages/client-runtime/src/authorization/service.ts +++ b/packages/client-runtime/src/authorization/service.ts @@ -33,7 +33,6 @@ import * as Effect from "effect/Effect"; import * as Fiber from "effect/Fiber"; import * as Layer from "effect/Layer"; import * as Option from "effect/Option"; -import * as Ref from "effect/Ref"; import * as Result from "effect/Result"; import * as Scope from "effect/Scope"; import * as Semaphore from "effect/Semaphore"; @@ -50,6 +49,7 @@ export interface AuthorizedRemoteEnvironment { readonly httpBaseUrl: string; readonly socketUrl: string; readonly httpAuthorization: PreparedHttpAuthorization; + readonly descriptor?: ExecutionEnvironmentDescriptor; } export interface AuthorizedRemoteHttpEnvironment { @@ -86,7 +86,6 @@ export class RemoteEnvironmentAuthorization extends Context.Service< >()("@t3tools/client-runtime/authorization/service/RemoteEnvironmentAuthorization") {} const CACHED_ENDPOINT_SOCKET_TIMEOUT_MS = 3_000; -const BEARER_DESCRIPTOR_CACHE_TTL_MS = 10_000; const DPOP_AUTHORIZATION_TIMEOUT_MS = 30_000; function mapDpopSocketError(error: RemoteEnvironmentAuthError | ConnectionAttemptError) { @@ -127,16 +126,6 @@ export const make = Effect.gen(function* () { readonly token: Effect.Effect; } >(); - const bearerDescriptors = yield* Ref.make< - ReadonlyMap< - EnvironmentId, - { - readonly httpBaseUrl: string; - readonly descriptor: ExecutionEnvironmentDescriptor; - readonly validatedAtEpochMs: number; - } - > - >(new Map()); /** * A saved LAN address can belong to another machine on a different network. @@ -147,33 +136,15 @@ export const make = Effect.gen(function* () { readonly httpBaseUrl: string; readonly connectionMethod: ClientConnectionMethod; }) { - const now = yield* Clock.currentTimeMillis; - const cachedDescriptor = (yield* Ref.get(bearerDescriptors)).get(input.expectedEnvironmentId); - const canReuseDescriptor = - cachedDescriptor?.httpBaseUrl === input.httpBaseUrl && - cachedDescriptor.validatedAtEpochMs + BEARER_DESCRIPTOR_CACHE_TTL_MS > now; - const descriptor = canReuseDescriptor - ? cachedDescriptor.descriptor - : yield* fetchDescriptor(input.httpBaseUrl, input.connectionMethod).pipe( - Effect.provideService(HttpClient.HttpClient, httpClient), - ); + const descriptor = yield* fetchDescriptor(input.httpBaseUrl, input.connectionMethod).pipe( + Effect.provideService(HttpClient.HttpClient, httpClient), + ); if (descriptor.environmentId !== input.expectedEnvironmentId) { return yield* environmentMismatchError({ expected: input.expectedEnvironmentId, actual: descriptor.environmentId, }); } - if (!canReuseDescriptor) { - yield* Ref.update(bearerDescriptors, (current) => { - const next = new Map(current); - next.set(input.expectedEnvironmentId, { - httpBaseUrl: input.httpBaseUrl, - descriptor, - validatedAtEpochMs: now, - }); - return next; - }); - } return descriptor; }, ); @@ -204,6 +175,7 @@ export const make = Effect.gen(function* () { label: descriptor.label, httpBaseUrl: input.httpBaseUrl, socketUrl, + descriptor, httpAuthorization: { _tag: "Bearer" as const, token: input.bearerToken, @@ -492,16 +464,18 @@ export const make = Effect.gen(function* () { input: Parameters[0], ) { const endpoint = input.directEndpoint; - if (endpoint !== undefined) { - yield* verifyDirectEndpoint({ - expectedEnvironmentId: input.expectedEnvironmentId, - httpBaseUrl: endpoint.httpBaseUrl, - connectionMethod: "direct", - }); - } + const descriptor = + endpoint === undefined + ? undefined + : yield* verifyDirectEndpoint({ + expectedEnvironmentId: input.expectedEnvironmentId, + httpBaseUrl: endpoint.httpBaseUrl, + connectionMethod: "direct", + }); const authorized = (token: TokenStore.RemoteDpopAccessToken, socketUrl: string) => ({ ...httpAuthorization(token, endpoint?.httpBaseUrl), socketUrl, + ...(descriptor === undefined ? {} : { descriptor }), }); let selected = yield* getDpopToken(input); if (selected.fromCache) { diff --git a/packages/client-runtime/src/connection/resolver.ts b/packages/client-runtime/src/connection/resolver.ts index 08908a64fd641..d3b22143d327a 100644 --- a/packages/client-runtime/src/connection/resolver.ts +++ b/packages/client-runtime/src/connection/resolver.ts @@ -171,11 +171,7 @@ const makeBearerBroker = Effect.fn("clientRuntime.connection.broker.makeBearer") connectionMethod: "direct", }); return { - environmentId: authorized.environmentId, - label: authorized.label, - httpBaseUrl: authorized.httpBaseUrl, - socketUrl: authorized.socketUrl, - httpAuthorization: authorized.httpAuthorization, + ...authorized, target, } satisfies PreparedConnection; }); @@ -190,11 +186,7 @@ const makeRelayBroker = Effect.fn("clientRuntime.connection.broker.makeRelay")(f expectedEnvironmentId: target.environmentId, }); return { - environmentId: authorized.environmentId, - label: authorized.label, - httpBaseUrl: authorized.httpBaseUrl, - socketUrl: authorized.socketUrl, - httpAuthorization: authorized.httpAuthorization, + ...authorized, target, } satisfies PreparedConnection; }, @@ -265,11 +257,7 @@ const makeSshBroker = Effect.fn("clientRuntime.connection.broker.makeSsh")(funct connectionMethod: "ssh", }); return { - environmentId: authorized.environmentId, - label: authorized.label, - httpBaseUrl: authorized.httpBaseUrl, - socketUrl: authorized.socketUrl, - httpAuthorization: authorized.httpAuthorization, + ...authorized, target, } satisfies PreparedConnection; }); @@ -291,7 +279,9 @@ export const make = Effect.gen(function* () { "connection.environment.id": target.environmentId, "connection.target.kind": target._tag, }); - const prepared = yield* (() => { + const authorized: PreparedConnection & { + readonly descriptor?: ExecutionEnvironmentDescriptor; + } = yield* (() => { switch (target._tag) { case "PrimaryConnectionTarget": return primary(target); @@ -303,12 +293,15 @@ export const make = Effect.gen(function* () { return ssh({ ...entry, target }); } })(); - const descriptor = yield* fetchRemoteEnvironmentDescriptor({ - httpBaseUrl: prepared.httpBaseUrl, - }).pipe( - Effect.mapError(mapRemoteEnvironmentError), - Effect.provideService(HttpClient.HttpClient, httpClient), - ); + const { descriptor: authorizedDescriptor, ...prepared } = authorized; + const descriptor = + authorizedDescriptor ?? + (yield* fetchRemoteEnvironmentDescriptor({ + httpBaseUrl: prepared.httpBaseUrl, + }).pipe( + Effect.mapError(mapRemoteEnvironmentError), + Effect.provideService(HttpClient.HttpClient, httpClient), + )); if (descriptor.environmentId !== target.environmentId) { return yield* environmentMismatchError({ expected: target.environmentId, diff --git a/packages/client-runtime/src/state/threadCommands.test.ts b/packages/client-runtime/src/state/threadCommands.test.ts index d61213c243ea8..2b1b1162d6bc9 100644 --- a/packages/client-runtime/src/state/threadCommands.test.ts +++ b/packages/client-runtime/src/state/threadCommands.test.ts @@ -3,6 +3,7 @@ import { CommandId, RuntimeRequestId, EnvironmentId, + MessageId, RunId, ORCHESTRATION_V2_WS_METHODS, ProjectId, @@ -86,6 +87,7 @@ const SNAPSHOT: OrchestrationV2ShellSnapshot = { }; const makeHarness = Effect.fn("TestThreadCommands.makeHarness")(function* (options?: { + readonly loadEarlier?: () => Effect.Effect; readonly reconcilePendingWork?: (threadId: ThreadId) => Effect.Effect; }) { const requests = yield* Queue.unbounded<{ @@ -126,7 +128,7 @@ const makeHarness = Effect.fn("TestThreadCommands.makeHarness")(function* (optio ThreadHistoryController.ThreadHistoryController.of({ register: () => Effect.die("unused"), unregister: () => Effect.void, - loadEarlier: () => Effect.die("unused"), + loadEarlier: () => options?.loadEarlier?.() ?? Effect.die("unused"), reconcilePendingWork: (_environmentId, threadId) => options?.reconcilePendingWork?.(threadId) ?? Queue.offer(reconciled, threadId), }), @@ -142,6 +144,59 @@ const makeHarness = Effect.fn("TestThreadCommands.makeHarness")(function* (optio return { registry, commands, snapshotAtom, visibleAtom, requests, reconciled }; }); +describe("history loading", () => { + it.effect("sends a message on the thread while earlier history is loading", () => + Effect.gen(function* () { + const historyStarted = yield* Deferred.make(); + const historyReply = + yield* Deferred.make(); + const h = yield* makeHarness({ + loadEarlier: () => + Deferred.succeed(historyStarted, undefined).pipe( + Effect.andThen(Deferred.await(historyReply)), + ), + }); + yield* Effect.addFinalizer(() => Deferred.succeed(historyReply, { _tag: "loaded" })); + const history = h.commands.loadEarlierHistory.run(h.registry, { + environmentId: ENVIRONMENT_ID, + input: { threadId: THREAD_ID }, + }); + yield* Deferred.await(historyStarted); + const send = h.commands.startTurn.run(h.registry, { + environmentId: ENVIRONMENT_ID, + input: { + threadId: THREAD_ID, + message: { + messageId: MessageId.make("message"), + role: "user", + text: "Send while history loads", + attachments: [], + }, + runtimeMode: "full-access", + interactionMode: "default", + dispatchMode: "start", + }, + }); + // Dispatched after Send on another thread, so it arrives first only if Send is queued. + const controlThreadId = ThreadId.make("control-thread"); + const control = h.commands.interruptTurn.run(h.registry, { + environmentId: ENVIRONMENT_ID, + input: { threadId: controlThreadId, runId: RunId.make("control-run") }, + }); + const first = yield* Queue.take(h.requests); + yield* Deferred.succeed(first.reply, { sequence: 1 }); + yield* Deferred.succeed(historyReply, { _tag: "loaded" }); + const second = yield* Queue.take(h.requests); + yield* Deferred.succeed(second.reply, { sequence: 2 }); + const results = yield* Effect.promise(() => Promise.all([history, send, control])); + + expect(first.command).toMatchObject({ type: "message.dispatch", threadId: THREAD_ID }); + expect(second.command).toMatchObject({ type: "run.interrupt", threadId: controlThreadId }); + expect(results.map((value) => value._tag)).toEqual(["Success", "Success", "Success"]); + }), + ); +}); + describe("interrupt turn", () => { // A stale cached item has nothing left to interrupt, so the server can reject Stop. it.effect("reconciles pending work when the server rejects the interrupt", () => diff --git a/packages/client-runtime/src/state/threadCommands.ts b/packages/client-runtime/src/state/threadCommands.ts index 293b69ef46434..6070638c57a66 100644 --- a/packages/client-runtime/src/state/threadCommands.ts +++ b/packages/client-runtime/src/state/threadCommands.ts @@ -351,7 +351,8 @@ export function createThreadEnvironmentAtoms( scheduler, concurrency: { mode: "serial", - key: ({ environmentId, input }) => JSON.stringify([environmentId, input.threadId]), + key: ({ environmentId, input }) => + `history:${JSON.stringify([environmentId, input.threadId])}`, }, }), uploadFeedback: createEnvironmentRpcCommand(runtime, {