diff --git a/apps/server/src/auth/RpcAuthorization.ts b/apps/server/src/auth/RpcAuthorization.ts index 880475227..0fbf3d2d5 100644 --- a/apps/server/src/auth/RpcAuthorization.ts +++ b/apps/server/src/auth/RpcAuthorization.ts @@ -57,6 +57,15 @@ export const RPC_REQUIRED_SCOPES = { [WS_METHODS.serverStartProviderLogin]: AuthOrchestrationOperateScope, [WS_METHODS.serverSubmitProviderLoginCode]: AuthOrchestrationOperateScope, [WS_METHODS.serverCancelProviderLogin]: AuthOrchestrationOperateScope, + [WS_METHODS.providerAuthStart]: AuthOrchestrationOperateScope, + [WS_METHODS.providerAuthComplete]: AuthOrchestrationOperateScope, + [WS_METHODS.providerAuthCancel]: AuthOrchestrationOperateScope, + [WS_METHODS.providerAuthLogout]: AuthOrchestrationOperateScope, + [WS_METHODS.providerAuthSubscribe]: AuthOrchestrationOperateScope, + [WS_METHODS.providerInstallStart]: AuthOrchestrationOperateScope, + [WS_METHODS.providerInstallCancel]: AuthOrchestrationOperateScope, + [WS_METHODS.providerInstallSubscribe]: AuthOrchestrationReadScope, + [WS_METHODS.providerInstallRemove]: AuthOrchestrationOperateScope, [WS_METHODS.serverUpdateServer]: AuthOrchestrationOperateScope, [WS_METHODS.serverUpdateServerWithProgress]: AuthOrchestrationOperateScope, [WS_METHODS.serverCommitDesktopUpdate]: AuthOrchestrationOperateScope, diff --git a/apps/server/src/server.test.ts b/apps/server/src/server.test.ts index b8f7d006b..730708f5a 100644 --- a/apps/server/src/server.test.ts +++ b/apps/server/src/server.test.ts @@ -1,5 +1,7 @@ import * as UsageLimitSources from "./usage/UsageLimitSources.ts"; import { ProviderInstanceRegistry } from "./provider/Services/ProviderInstanceRegistry.ts"; +import { ProviderAuthService } from "./provider/Services/ProviderAuthService.ts"; +import { AntigravityInstallation } from "./provider/AntigravityInstallation.ts"; import * as NodeHttpServer from "@effect/platform-node/NodeHttpServer"; import * as NodeSocket from "@effect/platform-node/NodeSocket"; import * as NodeServices from "@effect/platform-node/NodeServices"; @@ -34,6 +36,8 @@ import { ProjectId, ProviderDriverKind, ProviderInstanceId, + type ProviderAuthState, + type ProviderInstallState, UsageLimitSourceId, ServerProviderMutationBusyError, ResolvedKeybindingRule, @@ -440,6 +444,27 @@ const makeBrowserOtlpPayload = (spanName: string) => return JSON.parse(request.body) as OtlpTracer.TraceData; }); +const IDLE_PROVIDER_AUTH_STATE: ProviderAuthState = { + instanceId: ProviderInstanceId.make("antigravity"), + phase: "idle", + flowId: null, + authorizationUrl: null, + expiresAt: null, + message: null, +}; + +const IDLE_PROVIDER_INSTALL_STATE: ProviderInstallState = { + driver: ProviderDriverKind.make("antigravity"), + operationId: null, + phase: "idle", + downloadedBytes: 0, + totalBytes: null, + version: null, + installedVersion: null, + canRemove: false, + message: null, +}; + const buildAppUnderTest = (options?: { onPairingChangesSubscribed?: Effect.Effect; config?: Partial; @@ -447,6 +472,8 @@ const buildAppUnderTest = (options?: { keybindings?: Partial; usageLimitSources?: Partial; providerInstances?: Partial; + providerAuth?: Partial; + antigravityInstallation?: Partial; environmentTheme?: Partial; providerRegistry?: Partial; providerService?: Partial; @@ -990,6 +1017,30 @@ const buildAppUnderTest = (options?: { ...options?.layers?.serverLifecycleEvents, }), ), + Layer.provide( + Layer.mock(ProviderAuthService)({ + start: () => Effect.succeed(IDLE_PROVIDER_AUTH_STATE), + complete: () => Effect.succeed(IDLE_PROVIDER_AUTH_STATE), + cancel: () => Effect.succeed(IDLE_PROVIDER_AUTH_STATE), + logout: () => Effect.succeed(IDLE_PROVIDER_AUTH_STATE), + subscribe: () => Stream.make(IDLE_PROVIDER_AUTH_STATE), + tryHandlePromptCommand: () => Effect.succeed(false), + ...options?.layers?.providerAuth, + }), + ), + Layer.provide( + Layer.mock(AntigravityInstallation)({ + managedDirectory: "/tmp/antigravity", + resolve: () => Effect.die("resolve is not stubbed"), + acquire: () => Effect.die("acquire is not stubbed"), + start: Effect.succeed(IDLE_PROVIDER_INSTALL_STATE), + cancel: () => Effect.succeed(IDLE_PROVIDER_INSTALL_STATE), + state: Effect.succeed(IDLE_PROVIDER_INSTALL_STATE), + changes: Stream.make(IDLE_PROVIDER_INSTALL_STATE), + remove: () => Effect.void, + ...options?.layers?.antigravityInstallation, + }), + ), Layer.provide( Layer.mock(ServerRuntimeStartup.ServerRuntimeStartup)({ awaitCommandReady: Effect.void, diff --git a/apps/server/src/server.ts b/apps/server/src/server.ts index f75f0389e..1ca8708d9 100644 --- a/apps/server/src/server.ts +++ b/apps/server/src/server.ts @@ -37,6 +37,7 @@ import * as PrimeAgentRecoveryLedger from "./provider/prime/PrimeAgentRecoveryLe import { ProviderAdapterRegistryLive } from "./provider/Layers/ProviderAdapterRegistry.ts"; import * as ModelManifest from "./provider/ModelManifest.ts"; import { AntigravityInstallation } from "./provider/AntigravityInstallation.ts"; +import { ProviderAuthServiceLive } from "./provider/Layers/ProviderAuthService.ts"; import * as ProviderEventLoggers from "./provider/Layers/ProviderEventLoggers.ts"; import { ProviderServiceLive } from "./provider/Layers/ProviderService.ts"; import { ProviderSessionReaperLive } from "./provider/Layers/ProviderSessionReaper.ts"; @@ -298,10 +299,17 @@ const ProviderSessionDirectoryLayerLive = ProviderSessionDirectoryLive.pipe( // `create()`; `ProviderEventLoggers.layer` owns the shared native/canonical // NDJSON writers and is provided at the outer runtime layer so both // `ProviderService` and the per-instance drivers read the same logger pair. -const ProviderLayerLive = ProviderServiceLive.pipe( - Layer.provide(ProviderAdapterRegistryLive), - Layer.provideMerge(ProviderSessionDirectoryLayerLive), - Layer.provideMerge(RollbackSagaRepositoryLive), +// ProviderAuthService backs the provider setup/auth RPCs and reads the provider +// service and session directory, so it is layered on top of them rather than +// provided to them. +const ProviderLayerLive = ProviderAuthServiceLive.pipe( + Layer.provideMerge( + ProviderServiceLive.pipe( + Layer.provide(ProviderAdapterRegistryLive), + Layer.provideMerge(ProviderSessionDirectoryLayerLive), + Layer.provideMerge(RollbackSagaRepositoryLive), + ), + ), ); const PersistenceLayerLive = PrimeAgentRecoveryLedger.layer.pipe( diff --git a/apps/server/src/ws.ts b/apps/server/src/ws.ts index f8f1c3136..7e69a4637 100644 --- a/apps/server/src/ws.ts +++ b/apps/server/src/ws.ts @@ -57,6 +57,7 @@ import { ProjectWriteFileError, ProviderUploadFeedbackError, ProviderDriverKind, + ProviderSetupError, RelayClientInstallFailedError, type RelayClientInstallProgressEvent, ServerSelfUpdateError, @@ -126,6 +127,8 @@ import { ProviderLoginCoordinatorLive, } from "./provider/providerLoginCoordinator.ts"; import { ProviderLoginSessionsLive } from "./provider/ProviderLoginSessions.ts"; +import { ProviderAuthService } from "./provider/Services/ProviderAuthService.ts"; +import { makeProviderInstallation } from "./provider/providerInstallation.ts"; import * as ServerSelfUpdate from "./cloud/selfUpdate.ts"; import * as ServerLifecycleEvents from "./serverLifecycleEvents.ts"; import * as ServerRuntimeStartup from "./serverRuntimeStartup.ts"; @@ -526,6 +529,8 @@ const makeWsRpcLayer = ( const providerMaintenanceRunner = yield* ProviderMaintenanceRunner.ProviderMaintenanceRunner; const primeManagedMaintenance = yield* PrimeManagedMaintenance.PrimeManagedMaintenance; const providerLogin = yield* ProviderLoginCoordinator; + const providerAuth = yield* ProviderAuthService; + const providerInstallation = yield* makeProviderInstallation(); const serverUpdate = yield* ServerSelfUpdate.ServerSelfUpdate; const config = yield* ServerConfig.ServerConfig; const lifecycleEvents = yield* ServerLifecycleEvents.ServerLifecycleEvents; @@ -1793,20 +1798,47 @@ const makeWsRpcLayer = ( [WS_METHODS.serverRefreshProviders]: (input) => observeRpcEffect( WS_METHODS.serverRefreshProviders, - (input.cwd !== undefined && input.instanceId !== undefined - ? providerRegistry.refreshWorkspaceSnapshot({ - instanceId: input.instanceId, - cwd: input.cwd, - }) - : input.instanceId !== undefined - ? providerRegistry.refreshInstance(input.instanceId) - : providerRegistry.refresh() - ).pipe( - Effect.tap(() => - input.instanceId === undefined ? usageLimitSources.refresh : Effect.void, - ), - Effect.map((providers) => ({ providers })), - ), + Effect.gen(function* () { + let providers = yield* input.cwd !== undefined && input.instanceId !== undefined + ? providerRegistry.refreshWorkspaceSnapshot({ + instanceId: input.instanceId, + cwd: input.cwd, + }) + : input.instanceId !== undefined + ? providerRegistry.refreshInstance(input.instanceId) + : providerRegistry.refresh(); + if (input.instanceId === undefined) { + yield* usageLimitSources.refresh; + } + if (input.refreshModels) { + const instances = yield* providerInstances.listInstances; + for (const instance of instances) { + if ( + !instance.refreshModels || + (input.instanceId !== undefined && input.instanceId !== instance.instanceId) || + !providers.some( + (provider) => + provider.instanceId === instance.instanceId && + provider.enabled && + provider.installed, + ) + ) + continue; + yield* instance.refreshModels().pipe( + Effect.mapError( + (error) => + new ProviderSetupError({ + instanceId: instance.instanceId, + operation: "refresh-models", + detail: error.detail, + }), + ), + ); + providers = yield* providerRegistry.refreshInstance(instance.instanceId); + } + } + return { providers }; + }), { "rpc.aggregate": "server" }, ), [WS_METHODS.providerAskSessionSideQuestion]: (input) => @@ -2129,6 +2161,52 @@ const makeWsRpcLayer = ( observeRpcEffect(WS_METHODS.serverCancelProviderLogin, providerLogin.cancel(input), { "rpc.aggregate": "server", }), + [WS_METHODS.providerAuthStart]: (input) => + observeRpcEffect( + WS_METHODS.providerAuthStart, + providerAuth.start(input, currentSessionId), + { "rpc.aggregate": "provider" }, + ), + [WS_METHODS.providerAuthComplete]: (input) => + observeRpcEffect( + WS_METHODS.providerAuthComplete, + providerAuth.complete(input, currentSessionId), + { "rpc.aggregate": "provider" }, + ), + [WS_METHODS.providerAuthCancel]: (input) => + observeRpcEffect( + WS_METHODS.providerAuthCancel, + providerAuth.cancel(input, currentSessionId), + { "rpc.aggregate": "provider" }, + ), + [WS_METHODS.providerAuthLogout]: (input) => + observeRpcEffect(WS_METHODS.providerAuthLogout, providerAuth.logout(input), { + "rpc.aggregate": "provider", + }), + [WS_METHODS.providerAuthSubscribe]: (input) => + observeRpcStream( + WS_METHODS.providerAuthSubscribe, + providerAuth.subscribe(input, currentSessionId), + { "rpc.aggregate": "provider" }, + ), + [WS_METHODS.providerInstallStart]: (input) => + observeRpcEffect(WS_METHODS.providerInstallStart, providerInstallation.start(input), { + "rpc.aggregate": "provider", + }), + [WS_METHODS.providerInstallCancel]: (input) => + observeRpcEffect(WS_METHODS.providerInstallCancel, providerInstallation.cancel(input), { + "rpc.aggregate": "provider", + }), + [WS_METHODS.providerInstallSubscribe]: (input) => + observeRpcStream( + WS_METHODS.providerInstallSubscribe, + providerInstallation.subscribe(input), + { "rpc.aggregate": "provider" }, + ), + [WS_METHODS.providerInstallRemove]: (input) => + observeRpcEffect(WS_METHODS.providerInstallRemove, providerInstallation.remove(input), { + "rpc.aggregate": "provider", + }), [WS_METHODS.serverUpdateServer]: (input) => observeRpcEffect(WS_METHODS.serverUpdateServer, serverUpdate.update(input), { "rpc.aggregate": "server", diff --git a/packages/client-runtime/src/rpc/client.ts b/packages/client-runtime/src/rpc/client.ts index b176c32f2..be32fe8a9 100644 --- a/packages/client-runtime/src/rpc/client.ts +++ b/packages/client-runtime/src/rpc/client.ts @@ -44,6 +44,8 @@ export type EnvironmentRpcTag = keyof WsRpcProtocolClient & string; type RpcMethod = WsRpcProtocolClient[TTag]; export type EnvironmentSubscriptionRpcTag = + | typeof WS_METHODS.providerAuthSubscribe + | typeof WS_METHODS.providerInstallSubscribe | typeof ORCHESTRATION_WS_METHODS.subscribeShell | typeof ORCHESTRATION_WS_METHODS.subscribeThread | typeof WS_METHODS.subscribeAuthAccess diff --git a/packages/client-runtime/src/state/server.ts b/packages/client-runtime/src/state/server.ts index d4c038ab1..b2d8d3f19 100644 --- a/packages/client-runtime/src/state/server.ts +++ b/packages/client-runtime/src/state/server.ts @@ -941,6 +941,52 @@ export function createServerEnvironmentAtoms( updateStateAtom, settingsValueAtom, providersValueAtom, + providerAuthState: createEnvironmentRpcSubscriptionAtomFamily(runtime, { + label: "environment-data:provider:auth-state", + tag: WS_METHODS.providerAuthSubscribe, + idleTtlMs: 0, + }), + startProviderAuth: createEnvironmentRpcCommand(runtime, { + label: "environment-data:provider:auth-start", + tag: WS_METHODS.providerAuthStart, + concurrency: { + mode: "singleFlight", + key: ({ environmentId, input }) => JSON.stringify([environmentId, input.instanceId]), + }, + }), + completeProviderAuth: createEnvironmentRpcCommand(runtime, { + label: "environment-data:provider:auth-complete", + tag: WS_METHODS.providerAuthComplete, + }), + cancelProviderAuth: createEnvironmentRpcCommand(runtime, { + label: "environment-data:provider:auth-cancel", + tag: WS_METHODS.providerAuthCancel, + }), + logoutProviderAuth: createEnvironmentRpcCommand(runtime, { + label: "environment-data:provider:auth-logout", + tag: WS_METHODS.providerAuthLogout, + }), + providerInstallState: createEnvironmentRpcSubscriptionAtomFamily(runtime, { + label: "environment-data:provider:install-state", + tag: WS_METHODS.providerInstallSubscribe, + idleTtlMs: 0, + }), + startProviderInstall: createEnvironmentRpcCommand(runtime, { + label: "environment-data:provider:install-start", + tag: WS_METHODS.providerInstallStart, + concurrency: { + mode: "singleFlight", + key: ({ environmentId }) => environmentId, + }, + }), + cancelProviderInstall: createEnvironmentRpcCommand(runtime, { + label: "environment-data:provider:install-cancel", + tag: WS_METHODS.providerInstallCancel, + }), + removeProviderInstallation: createEnvironmentRpcCommand(runtime, { + label: "environment-data:provider:install-remove", + tag: WS_METHODS.providerInstallRemove, + }), traceDiagnostics: createEnvironmentRpcQueryAtomFamily(runtime, { label: "environment-data:server:trace-diagnostics", tag: WS_METHODS.serverGetTraceDiagnostics, @@ -999,7 +1045,12 @@ export function createServerEnvironmentAtoms( concurrency: { mode: "singleFlight", key: ({ environmentId, input }) => - JSON.stringify([environmentId, input.instanceId ?? null, input.cwd ?? null]), + JSON.stringify([ + environmentId, + input.instanceId ?? null, + input.cwd ?? null, + input.refreshModels ?? false, + ]), }, }), updateProvider: createEnvironmentRpcCommand(runtime, { diff --git a/packages/contracts/src/rpc.ts b/packages/contracts/src/rpc.ts index ed303a110..52e050ea5 100644 --- a/packages/contracts/src/rpc.ts +++ b/packages/contracts/src/rpc.ts @@ -8,6 +8,15 @@ import * as Schema from "effect/Schema"; import * as Rpc from "effect/unstable/rpc/Rpc"; import * as RpcGroup from "effect/unstable/rpc/RpcGroup"; import { NonNegativeInt, TrimmedNonEmptyString } from "./baseSchemas.ts"; +import { + ProviderAuthCancelInput, + ProviderAuthCompleteInput, + ProviderAuthState, + ProviderInstallCancelInput, + ProviderInstallState, + ProviderSetupError, + ProviderSetupInput, +} from "./providerSetup.ts"; import { ExternalLauncherError, LaunchEditorInput } from "./editor.ts"; import { @@ -302,6 +311,15 @@ export const WS_METHODS = { // Provider methods providerUploadFeedback: "provider.uploadFeedback", + providerAuthStart: "provider.auth.start", + providerAuthComplete: "provider.auth.complete", + providerAuthCancel: "provider.auth.cancel", + providerAuthLogout: "provider.auth.logout", + providerAuthSubscribe: "provider.auth.subscribe", + providerInstallStart: "provider.install.start", + providerInstallCancel: "provider.install.cancel", + providerInstallSubscribe: "provider.install.subscribe", + providerInstallRemove: "provider.install.remove", // VCS methods vcsPull: "vcs.pull", @@ -476,9 +494,11 @@ const WsServerRefreshProvidersRpc = Rpc.make(WS_METHODS.serverRefreshProviders, */ instanceId: Schema.optional(ProviderInstanceId), cwd: Schema.optional(TrimmedNonEmptyString), + /** Explicit user request. Background status refreshes must not open agent sessions. */ + refreshModels: Schema.optional(Schema.Boolean), }), success: ServerProviderUpdatedPayload, - error: EnvironmentAuthorizationError, + error: Schema.Union([EnvironmentAuthorizationError, ProviderSetupError]), }); const WsProviderReloadSessionResourcesRpc = Rpc.make(WS_METHODS.providerReloadSessionResources, { @@ -673,6 +693,64 @@ const WsServerCancelProviderLoginRpc = Rpc.make(WS_METHODS.serverCancelProviderL error: EnvironmentAuthorizationError, }); +const ProviderSetupRpcError = Schema.Union([ProviderSetupError, EnvironmentAuthorizationError]); + +const WsProviderAuthStartRpc = Rpc.make(WS_METHODS.providerAuthStart, { + payload: ProviderSetupInput, + success: ProviderAuthState, + error: ProviderSetupRpcError, +}); + +const WsProviderAuthCompleteRpc = Rpc.make(WS_METHODS.providerAuthComplete, { + payload: ProviderAuthCompleteInput, + success: ProviderAuthState, + error: ProviderSetupRpcError, +}); + +const WsProviderAuthCancelRpc = Rpc.make(WS_METHODS.providerAuthCancel, { + payload: ProviderAuthCancelInput, + success: ProviderAuthState, + error: ProviderSetupRpcError, +}); + +const WsProviderAuthLogoutRpc = Rpc.make(WS_METHODS.providerAuthLogout, { + payload: ProviderSetupInput, + success: ProviderAuthState, + error: ProviderSetupRpcError, +}); + +const WsProviderAuthSubscribeRpc = Rpc.make(WS_METHODS.providerAuthSubscribe, { + payload: ProviderSetupInput, + success: ProviderAuthState, + error: ProviderSetupRpcError, + stream: true, +}); + +const WsProviderInstallStartRpc = Rpc.make(WS_METHODS.providerInstallStart, { + payload: ProviderSetupInput, + success: ProviderInstallState, + error: ProviderSetupRpcError, +}); + +const WsProviderInstallCancelRpc = Rpc.make(WS_METHODS.providerInstallCancel, { + payload: ProviderInstallCancelInput, + success: ProviderInstallState, + error: ProviderSetupRpcError, +}); + +const WsProviderInstallSubscribeRpc = Rpc.make(WS_METHODS.providerInstallSubscribe, { + payload: ProviderSetupInput, + success: ProviderInstallState, + error: ProviderSetupRpcError, + stream: true, +}); + +const WsProviderInstallRemoveRpc = Rpc.make(WS_METHODS.providerInstallRemove, { + payload: ProviderSetupInput, + success: ProviderInstallState, + error: ProviderSetupRpcError, +}); + const WsServerUpdateServerRpc = Rpc.make(WS_METHODS.serverUpdateServer, { payload: ServerSelfUpdateInput, success: ServerSelfUpdateResult, @@ -1381,6 +1459,15 @@ export const WsRpcGroup = RpcGroup.make( WsServerStartProviderLoginRpc, WsServerSubmitProviderLoginCodeRpc, WsServerCancelProviderLoginRpc, + WsProviderAuthStartRpc, + WsProviderAuthCompleteRpc, + WsProviderAuthCancelRpc, + WsProviderAuthLogoutRpc, + WsProviderAuthSubscribeRpc, + WsProviderInstallStartRpc, + WsProviderInstallCancelRpc, + WsProviderInstallSubscribeRpc, + WsProviderInstallRemoveRpc, WsServerUpdateServerRpc, WsServerUpdateServerWithProgressRpc, WsServerCommitDesktopUpdateRpc,