Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
59 changes: 58 additions & 1 deletion services/agents-login/src/worker/session.ts
Original file line number Diff line number Diff line change
Expand Up @@ -42,6 +42,8 @@ function defaultCommand(provider: Provider): { file: string; args: string[] } {
const MAX_BUFFER = 256 * 1024
const DEFAULT_REDIRECT_TIMEOUT_MS = 60_000
const CLAUDE_CREDENTIAL_POLL_MS = 250
const CLAUDE_SUBMIT_ENTER_DELAY_MS = 120
const CLAUDE_SUBMIT_RETRY_DELAY_MS = 3_000
const LOG_LINE_LIMIT = 512
const FAILURE_TAIL_LIMIT = 2_000

Expand All @@ -60,6 +62,8 @@ export class LoginSession {
private timer?: NodeJS.Timeout
private redirectTimer?: NodeJS.Timeout
private credentialProbeTimer?: NodeJS.Timeout
private claudeSubmitEnterTimer?: NodeJS.Timeout
private claudeSubmitRetryTimer?: NodeJS.Timeout
private redirectSubmitted = false
private pendingClaudeCode?: string
private updatedBy = 'unknown'
Expand Down Expand Up @@ -261,7 +265,8 @@ export class LoginSession {
provider: this.provider,
codeBytes: code.length,
})
this.proc.write(`${code}\r`)
this.proc.write(code)
this.scheduleClaudeSubmitEnter()
// Stay in awaiting_url so the success line emitted by the CLI after the
// paste-back still drives finalize(); only finalize() flips to finalizing.
this.message = 'Submitted the authorization code; completing login…'
Expand Down Expand Up @@ -343,6 +348,57 @@ export class LoginSession {
return { ok: true }
}

private scheduleClaudeSubmitEnter(): void {
this.clearClaudeSubmitTimers()
this.claudeSubmitEnterTimer = setTimeout(() => {
this.claudeSubmitEnterTimer = undefined
this.writeClaudeSubmitEnter('initial delayed Enter')
this.scheduleClaudeSubmitRetry()
}, CLAUDE_SUBMIT_ENTER_DELAY_MS)
if (typeof this.claudeSubmitEnterTimer.unref === 'function') {
this.claudeSubmitEnterTimer.unref()
}
}

private scheduleClaudeSubmitRetry(): void {
if (isTerminal(this.phase) || this.phase === 'finalizing') {
return
}
this.claudeSubmitRetryTimer = setTimeout(() => {
this.claudeSubmitRetryTimer = undefined
if (this.phase !== 'awaiting_url') {
return
}
this.writeClaudeSubmitEnter('retry Enter after delayed submit')
}, CLAUDE_SUBMIT_RETRY_DELAY_MS)
if (typeof this.claudeSubmitRetryTimer.unref === 'function') {
this.claudeSubmitRetryTimer.unref()
}
}

private writeClaudeSubmitEnter(reason: string): void {
if (!this.proc || !this.redirectSubmitted || isTerminal(this.phase) || this.phase === 'finalizing') {
return
}
this.deps.logger.info('Claude authorization code submit Enter written to PTY', {
sessionId: this.id,
provider: this.provider,
reason,
})
this.proc.write('\r')
}

private clearClaudeSubmitTimers(): void {
if (this.claudeSubmitEnterTimer) {
clearTimeout(this.claudeSubmitEnterTimer)
this.claudeSubmitEnterTimer = undefined
}
if (this.claudeSubmitRetryTimer) {
clearTimeout(this.claudeSubmitRetryTimer)
this.claudeSubmitRetryTimer = undefined
}
}

private startRedirectTimeout(): void {
const timeoutMs = this.deps.redirectTimeoutMs ?? DEFAULT_REDIRECT_TIMEOUT_MS
if (this.redirectTimer) {
Expand Down Expand Up @@ -415,6 +471,7 @@ export class LoginSession {
clearTimeout(this.credentialProbeTimer)
this.credentialProbeTimer = undefined
}
this.clearClaudeSubmitTimers()
}

private async finalize(updatedBy: string, claudeTokenOverride?: string, reason = 'completion detected'): Promise<void> {
Expand Down
77 changes: 67 additions & 10 deletions services/agents-login/test/session.test.ts
Original file line number Diff line number Diff line change
Expand Up @@ -13,6 +13,17 @@ async function tick(times = 8): Promise<void> {
}
}

async function waitForWrites(writes: () => string[], count: number, timeoutMs = 500): Promise<string[]> {
const start = Date.now()
for (;;) {
const current = writes()
if (current.length >= count || Date.now() - start > timeoutMs) {
return current
}
await new Promise((r) => setTimeout(r, 10))
}
}

interface PostedCredential {
userId: string
provider: 'CLAUDE' | 'CODEX'
Expand Down Expand Up @@ -115,10 +126,16 @@ describe('LoginSession state machine', () => {
{ type: 'emit', data: 'Visit https://claude.ai/oauth/authorize?code=1\r\nPaste code here if prompted >' },
{
type: 'expectStdin',
match: (i) => i === 'redirect-code-2\r',
match: (i) => i === 'redirect-code-2',
then: [
{ type: 'emit', data: `Login successful.\r\nYour OAuth token: ${token}\r\n` },
{ type: 'exit', code: 0 },
{
type: 'expectStdin',
match: (i) => i === '\r',
then: [
{ type: 'emit', data: `Login successful.\r\nYour OAuth token: ${token}\r\n` },
{ type: 'exit', code: 0 },
],
},
],
},
])
Expand All @@ -139,7 +156,36 @@ describe('LoginSession state machine', () => {
expect(sub.ok).toBe(true)

expect(await waitPhase(mgr, started.id, ['succeeded', 'failed'])).toBe('succeeded')
expect(instances[0].writes).toEqual(['redirect-code-2\r'])
expect(instances[0].writes).toEqual(['redirect-code-2', '\r'])
expect(agentsApi.posts).toEqual([
{ userId: 'alice', provider: 'CLAUDE', payload: { oauth_token: token } },
])
})

it('writes the Claude authorization code and submit Enter as distinct PTY writes before finalizing', async () => {
const token = 'sk-ant-oat01-SeparateEnterToken1234567890_-abcXYZ'
const { deps: d, instances } = deps(() => [
{ type: 'emit', data: 'Open https://claude.com/cai/oauth/authorize?code=true\r\nPaste code here if prompted >' },
{
type: 'expectStdin',
match: (i) => i === 'bulk-paste-sensitive-code',
then: [
{
type: 'expectStdin',
match: (i) => i === '\r',
then: [{ type: 'emit', data: `Your OAuth token: ${token}\r\n` }],
},
],
},
])
const mgr = new SessionManager(d)
const started = mgr.start('claude', 'alice')
await tick()

expect(mgr.submitRedirectUrl(started.id, 'bulk-paste-sensitive-code').ok).toBe(true)

expect(await waitPhase(mgr, started.id, ['succeeded', 'failed'], 1000)).toBe('succeeded')
expect(instances[0].writes).toEqual(['bulk-paste-sensitive-code', '\r'])
expect(agentsApi.posts).toEqual([
{ userId: 'alice', provider: 'CLAUDE', payload: { oauth_token: token } },
])
Expand Down Expand Up @@ -185,10 +231,11 @@ describe('LoginSession state machine', () => {
expect(mgr.submitRedirectUrl(started.id, 'other-code').ok).toBe(false)

dataCbs.forEach((cb) => cb('Paste code here if prompted >'))
await tick()
await waitForWrites(() => writes, 2)

expect(writes).toEqual(['queued-code-123\r'])
expect(writes).toEqual(['queued-code-123', '\r'])
expect(exitCbs.length).toBe(1)
expect(mgr.cancel(started.id).ok).toBe(true)
})

it('writes Claude authorization codes to the live PTY and fails if the handle is missing', async () => {
Expand All @@ -203,7 +250,8 @@ describe('LoginSession state machine', () => {
expect(
live.submitRedirectUrl('https://platform.claude.com/oauth/code/callback?code=live-code-123&state=state-1').ok,
).toBe(true)
expect(instances[0].writes).toEqual(['live-code-123\r'])
expect(await waitForWrites(() => instances[0].writes, 2)).toEqual(['live-code-123', '\r'])
live.cancel()

const logs: LogEntry[] = []
const { deps: missingDeps } = deps(
Expand Down Expand Up @@ -425,7 +473,8 @@ describe('LoginSession state machine', () => {
expect(mgr.submitRedirectUrl(started.id, ' ').ok).toBe(false)
// setup-token returns a bare code, not a URL — it must be accepted.
expect(mgr.submitRedirectUrl(started.id, 'ABCD-1234-token').ok).toBe(true)
expect(instances[0].writes).toEqual(['ABCD-1234-token\r'])
expect(await waitForWrites(() => instances[0].writes, 2)).toEqual(['ABCD-1234-token', '\r'])
expect(mgr.cancel(started.id).ok).toBe(true)
})

it('rejects redirect submission for the codex provider', async () => {
Expand Down Expand Up @@ -523,8 +572,16 @@ describe('LoginSession state machine', () => {
{ type: 'emit', data: 'Visit https://claude.ai/oauth/authorize?code=1\r\nPaste code here if prompted >' },
{
type: 'expectStdin',
match: (i) => i === 'ingest-failure-code\r',
then: [{ type: 'emit', data: 'Login successful.\r\nYour OAuth token: sk-ant-oat01-FailureToken1234567890\r\n' }],
match: (i) => i === 'ingest-failure-code',
then: [
{
type: 'expectStdin',
match: (i) => i === '\r',
then: [
{ type: 'emit', data: 'Login successful.\r\nYour OAuth token: sk-ant-oat01-FailureToken1234567890\r\n' },
],
},
],
},
])
const mgr = new SessionManager({
Expand Down