diff --git a/src/agents/plugins/codemie-code-hooks/__tests__/shell-hooks-source.test.ts b/src/agents/plugins/codemie-code-hooks/__tests__/shell-hooks-source.test.ts index 31b34a81a..d88901d2d 100644 --- a/src/agents/plugins/codemie-code-hooks/__tests__/shell-hooks-source.test.ts +++ b/src/agents/plugins/codemie-code-hooks/__tests__/shell-hooks-source.test.ts @@ -40,12 +40,18 @@ let capturePath: string; let modulePath: string; /** Load the plugin factory the way the OpenCode runtime would. */ -async function loadHooks(hookNames: string[]): Promise> { +async function loadHooks( + hookNames: string[], + commandOverrides: Record = {}, +): Promise> { const hooks: Record = {}; for (const name of hookNames) { // A shell command that appends whatever arrives on stdin, standing in for // the real `codemie hook` binary. - hooks[name] = [{ hooks: [{ type: 'command', command: `cat >> ${capturePath}` }] }]; + hooks[name] = [{ hooks: [{ + type: 'command', + command: commandOverrides[name] ?? `cat >> ${capturePath}`, + }] }]; } process.env.OPENCODE_HOOKS = JSON.stringify({ hooks }); @@ -121,6 +127,17 @@ describe('SHELL_HOOKS_PLUGIN_SOURCE', () => { expect(payloads[0].transcript_path).toBe(TRANSCRIPT); }); + it('propagates a blocking UserPromptSubmit hook result', async () => { + const hooks = await loadHooks(['UserPromptSubmit'], { UserPromptSubmit: 'exit 2' }); + + await expect( + hooks['chat.message']( + { sessionID: 'ses_blocked' }, + { parts: [{ type: 'text', text: 'hello' }] }, + ), + ).rejects.toThrow('Hook blocked execution'); + }); + it('collapses session.idle and session.status{idle} into a single Stop', async () => { const hooks = await loadHooks(['Stop']); diff --git a/src/agents/plugins/codemie-code-hooks/shell-hooks-source.ts b/src/agents/plugins/codemie-code-hooks/shell-hooks-source.ts index 995d95c8e..1260a886c 100644 --- a/src/agents/plugins/codemie-code-hooks/shell-hooks-source.ts +++ b/src/agents/plugins/codemie-code-hooks/shell-hooks-source.ts @@ -393,7 +393,9 @@ const ShellHooksPlugin: Plugin = async (_input) => ({ }, // UserPromptSubmit → chat.message. Marks the start of an active period; - // codemie hook forwards it to SessionStore.startActivityTracking. + // codemie hook forwards it to SessionStore.startActivityTracking. Unlike + // telemetry-only events, this hook must propagate exit code 2 so an auth gate + // can prevent the model request from being submitted. "chat.message": async (input: any, output: any) => { const sessionId = resolveSessionId(input); if (!sessionId) return; @@ -418,7 +420,12 @@ const ShellHooksPlugin: Plugin = async (_input) => ({ .join("\\n") : ""; - runCommands("UserPromptSubmit", sessionId, payload, commands, true); + const responses = runCommands("UserPromptSubmit", sessionId, payload, commands, false); + for (const response of responses) { + if (response.blocked) { + throw new Error(response.reason || "Hook blocked prompt submission"); + } + } }, // PreCompact → experimental.session.compacting (non-blocking) diff --git a/src/agents/plugins/opencode/__tests__/opencode-session-lifecycle.test.ts b/src/agents/plugins/opencode/__tests__/opencode-session-lifecycle.test.ts index 9a8b56c67..b666a7caf 100644 --- a/src/agents/plugins/opencode/__tests__/opencode-session-lifecycle.test.ts +++ b/src/agents/plugins/opencode/__tests__/opencode-session-lifecycle.test.ts @@ -103,6 +103,7 @@ describe('OpenCode beforeRun', () => { const hooks = JSON.parse(env.OPENCODE_HOOKS!).hooks; expect(Object.keys(hooks).sort()).toEqual(['Stop', 'UserPromptSubmit']); + expect(hooks.UserPromptSubmit[0].hooks[0].timeout).toBe(130); // Stop is detached: the plugin runs sync hooks with execSync, which would // otherwise stall OpenCode while the child re-parses SQLite. expect(hooks.Stop[0].hooks[0].async).toBe(true); diff --git a/src/agents/plugins/opencode/opencode.plugin.ts b/src/agents/plugins/opencode/opencode.plugin.ts index 3c6e95623..ccffda0f6 100644 --- a/src/agents/plugins/opencode/opencode.plugin.ts +++ b/src/agents/plugins/opencode/opencode.plugin.ts @@ -34,7 +34,9 @@ const OPENCODE_CLIENT_TYPE = 'codemie-opencode'; * child re-parses SQLite and rewrites JSONL. */ const DEFAULT_HOOKS: Record = { - UserPromptSubmit: [{ hooks: [{ type: 'command', command: 'codemie hook', timeout: 5 }] }], + // The auth gate may wait for the browser callback before it returns the + // blocking exit status. Keep the timeout above its 120-second auth window. + UserPromptSubmit: [{ hooks: [{ type: 'command', command: 'codemie hook', timeout: 130 }] }], Stop: [{ hooks: [{ type: 'command', command: 'codemie hook', timeout: 10, async: true }] }], }; diff --git a/src/providers/plugins/sso/proxy/plugins/__tests__/sso-auth.plugin.test.ts b/src/providers/plugins/sso/proxy/plugins/__tests__/sso-auth.plugin.test.ts new file mode 100644 index 000000000..cf96a2dec --- /dev/null +++ b/src/providers/plugins/sso/proxy/plugins/__tests__/sso-auth.plugin.test.ts @@ -0,0 +1,129 @@ +/** + * SSO proxy response handling tests. + * @group unit + */ + +import { Readable } from 'node:stream'; +import { describe, expect, it, vi } from 'vitest'; +import type { IncomingMessage } from 'node:http'; +import type { PluginContext, UpstreamResponseTools } from '../types.js'; +import type { ProxyContext } from '../../proxy-types.js'; +import type { SSOCredentials } from '../../../../../core/types.js'; +import { SSOAuthPlugin } from '../sso-auth.plugin.js'; + +const { markAnalyticsAuthInvalid } = vi.hoisted(() => ({ + markAnalyticsAuthInvalid: vi.fn().mockResolvedValue(undefined), +})); + +vi.mock('../../../../../../utils/analytics-auth-status.js', () => ({ + markAnalyticsAuthInvalid, +})); + +vi.mock('../../../../../../utils/logger.js', () => ({ + logger: { + debug: vi.fn(), + warn: vi.fn(), + }, +})); + +function makeResponse(statusCode: number, contentType: string): IncomingMessage { + const response = Readable.from([Buffer.from('login')]) as unknown as IncomingMessage; + response.statusCode = statusCode; + response.statusMessage = statusCode === 200 ? 'OK' : 'Unauthorized'; + response.headers = { 'content-type': contentType }; + return response; +} + +function makeTools(): UpstreamResponseTools { + return { + readBody: vi.fn(async () => Buffer.alloc(0)), + retry: vi.fn(), + fromBuffer: vi.fn((source: IncomingMessage, body: Buffer) => { + const response = Readable.from(body) as unknown as IncomingMessage; + response.statusCode = source.statusCode; + response.statusMessage = source.statusMessage; + response.headers = { ...source.headers }; + return response; + }), + }; +} + +function makeContext(): ProxyContext { + return { + requestId: 'request-1', + sessionId: 'session-1', + agentName: 'opencode', + method: 'POST', + url: '/responses', + headers: {}, + requestBody: null, + requestStartTime: Date.now(), + targetUrl: 'https://api.example.com/responses', + metadata: {}, + }; +} + +function makeCredentials(): SSOCredentials { + return { + cookies: { session: 'expired' }, + apiUrl: 'https://api.example.com', + }; +} + +function makePluginContext(credentials: SSOCredentials): PluginContext { + return { + config: { targetApiUrl: credentials.apiUrl, provider: 'ai-run-sso' }, + logger: { + debug: vi.fn(), + warn: vi.fn(), + } as never, + credentials, + }; +} + +async function readBody(response: IncomingMessage): Promise { + const chunks: Buffer[] = []; + for await (const chunk of response) { + chunks.push(Buffer.from(chunk)); + } + return Buffer.concat(chunks).toString('utf8'); +} + +describe('SSOAuthPlugin upstream response handling', () => { + it.each([ + [401, 'application/json'], + [403, 'application/json'], + [200, 'text/html; charset=utf-8'], + ])('converts HTTP %s authentication responses into a structured 401', async (statusCode, contentType) => { + const plugin = new SSOAuthPlugin(); + const interceptor = await plugin.createInterceptor(makePluginContext(makeCredentials())); + const response = await interceptor.onUpstreamResponse!( + makeContext(), + makeResponse(statusCode, contentType), + makeTools(), + ); + + expect(response.statusCode).toBe(401); + expect(response.headers['content-type']).toBe('application/json'); + expect(JSON.parse(await readBody(response))).toMatchObject({ + error: { + type: 'authentication_error', + code: 'AUTH_FAILED', + }, + }); + expect(markAnalyticsAuthInvalid).toHaveBeenCalledWith( + expect.any(String), + 'https://api.example.com', + ); + }); + + it('leaves a valid JSON model response unchanged', async () => { + const plugin = new SSOAuthPlugin(); + const interceptor = await plugin.createInterceptor(makePluginContext(makeCredentials())); + const response = makeResponse(200, 'application/json'); + + await expect( + interceptor.onUpstreamResponse!(makeContext(), response, makeTools()), + ).resolves.toBe(response); + }); +}); diff --git a/src/providers/plugins/sso/proxy/plugins/sso-auth.plugin.ts b/src/providers/plugins/sso/proxy/plugins/sso-auth.plugin.ts index 61ef326ba..38acb7957 100644 --- a/src/providers/plugins/sso/proxy/plugins/sso-auth.plugin.ts +++ b/src/providers/plugins/sso/proxy/plugins/sso-auth.plugin.ts @@ -6,11 +6,13 @@ * KISS: Simple interceptor, one clear purpose */ -import { ProxyPlugin, PluginContext, ProxyInterceptor } from './types.js'; +import { IncomingMessage } from 'http'; +import { ProxyPlugin, PluginContext, ProxyInterceptor, UpstreamResponseTools } from './types.js'; import { ProxyContext } from '../proxy-types.js'; import { SSOCredentials } from '../../../../core/types.js'; import { logger } from '../../../../../utils/logger.js'; import { AuthenticationError } from '../proxy-errors.js'; +import { markAnalyticsAuthInvalid } from '../../../../../utils/analytics-auth-status.js'; export class SSOAuthPlugin implements ProxyPlugin { id = '@codemie/proxy-sso-auth'; @@ -90,4 +92,52 @@ class SSOAuthInterceptor implements ProxyInterceptor { headerLength: cookieHeader.length }); } + + async onUpstreamResponse( + context: ProxyContext, + response: IncomingMessage, + tools: UpstreamResponseTools, + ): Promise { + const statusCode = response.statusCode || 0; + const contentType = String(response.headers['content-type'] || '').toLowerCase(); + const isHtmlLoginPage = statusCode >= 200 && statusCode < 300 && contentType.includes('text/html'); + const isUnauthorized = statusCode === 401 || statusCode === 403; + + if (!isHtmlLoginPage && !isUnauthorized) { + return response; + } + + const reason = isHtmlLoginPage + ? 'CodeMie SSO returned an HTML login page instead of an API response' + : `CodeMie SSO rejected the upstream request with HTTP ${statusCode}`; + + // Stop the invalid upstream body before replacing it with a small, stable + // JSON error. This also releases the pooled connection promptly. + response.resume(); + await markAnalyticsAuthInvalid(reason, this.credentials.apiUrl); + + const body = Buffer.from(JSON.stringify({ + error: { + type: 'authentication_error', + code: 'AUTH_FAILED', + message: 'CodeMie SSO session expired. Run `codemie profile login`, then restart the proxy.', + }, + })); + const replacement = tools.fromBuffer(response, body); + replacement.statusCode = 401; + replacement.statusMessage = 'Unauthorized'; + replacement.headers = { + 'content-type': 'application/json', + 'content-length': String(body.length), + }; + + logger.warn(`[${this.name}] Replaced upstream authentication response`, { + requestId: context.requestId, + statusCode, + contentType, + reason, + }); + + return replacement; + } }