Skip to content

CloudAgentNext - Apply balance-skip eligibility to command sendMessage follow-ups #4477

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
eshurakov merged 2 commits into from
Jul 9, 2026
Merged

CloudAgentNext - Apply balance-skip eligibility to command sendMessage follow-ups #4477

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
7e11036
fix(web): apply balance-skip eligibility to command sendMessage follo…
eshurakov Jul 9, 2026
22550c9
fix(web): fall back to balance-checked client when command-turn getSe…
eshurakov Jul 9, 2026
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
116 changes: 115 additions & 1 deletion apps/web/src/routers/cloud-agent-next-router.test.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -42,9 +42,12 @@ const mockGenerateCloudAgentAttachmentUploadUrl = jest.fn<
}) => Promise<{ signedUrl: string; key: string; expiresAt: string }>
>(() => Promise.resolve({ signedUrl: 'signed', key: 'key', expiresAt: 'expires' }));

const mockGetSession = jest.fn<(cloudAgentSessionId: string) => Promise<{ model?: string }>>();

const mockCreateCloudAgentNextClient = jest.fn(() => ({
prepareSession: mockPrepareSession,
sendMessage: mockSendMessage,
getSession: mockGetSession,
}));

const mockCreateCloudAgentNextClientForModel = jest.fn(
Expand Down Expand Up @@ -148,7 +151,9 @@ let createCaller: (ctx: { user: User }) => {
}>;
sendMessage: (input: {
cloudAgentSessionId: string;
payload: { type: 'prompt'; prompt: string; mode: string; model: string };
payload:
| { type: 'prompt'; prompt: string; mode: string; model: string }
| { type: 'command'; command: string; arguments: string };
attachments?: { path: string; files: string[] };
images?: { path: string; files: string[] };
}) => Promise<unknown>;
Expand Down Expand Up @@ -179,6 +184,11 @@ describe('cloudAgentNextRouter attachment forwarding', () => {
mockVerifyUserOwnsSessionV2ByCloudAgentId.mockResolvedValue({
kiloSessionId: 'ses_12345678901234567890123456',
});
mockComputeCloudAgentNextBalanceCheckEligibility.mockResolvedValue({
isFree: false,
hasUserByokAvailable: false,
});
mockGetSession.mockResolvedValue({ model: 'kilo/paid-model' });
});

it('denies a session the authenticated user does not own before calling the Worker', async () => {
Expand Down Expand Up @@ -228,6 +238,110 @@ describe('cloudAgentNextRouter attachment forwarding', () => {
expect(mockSendMessage).not.toHaveBeenCalledWith(expect.objectContaining({ images }));
});

it('routes free follow-up prompt models through the balance-skip client', async () => {
mockComputeCloudAgentNextBalanceCheckEligibility.mockResolvedValueOnce({
isFree: true,
hasUserByokAvailable: false,
});
const caller = createCaller({ user: { id: 'user-free', is_admin: false } as User });

await caller.sendMessage({
cloudAgentSessionId: 'agent_123',
payload: {
type: 'prompt',
prompt: 'Follow up on this',
mode: 'code',
model: 'kilo/free-model',
},
});

expect(mockComputeCloudAgentNextBalanceCheckEligibility).toHaveBeenCalledWith(
expect.objectContaining({ modelId: 'kilo/free-model' })
);
expect(mockCreateCloudAgentNextClientForModel).toHaveBeenCalledWith('cloud-agent-token', {
isFree: true,
hasUserByokAvailable: false,
});
expect(mockCreateCloudAgentNextClient).not.toHaveBeenCalled();
});

it('routes free follow-up command turns through the balance-skip client using the session model', async () => {
mockGetSession.mockResolvedValueOnce({ model: 'kilo/free-model' });
mockComputeCloudAgentNextBalanceCheckEligibility.mockResolvedValueOnce({
isFree: true,
hasUserByokAvailable: false,
});
const caller = createCaller({ user: { id: 'user-free', is_admin: false } as User });

await caller.sendMessage({
cloudAgentSessionId: 'agent_123',
payload: { type: 'command', command: 'review', arguments: '' },
});

expect(mockGetSession).toHaveBeenCalledWith('agent_123');
expect(mockComputeCloudAgentNextBalanceCheckEligibility).toHaveBeenCalledWith(
expect.objectContaining({ modelId: 'kilo/free-model' })
);
expect(mockCreateCloudAgentNextClientForModel).toHaveBeenCalledWith('cloud-agent-token', {
isFree: true,
hasUserByokAvailable: false,
});
});

it('keeps the balance check for command turns on paid sessions', async () => {
mockGetSession.mockResolvedValueOnce({ model: 'kilo/paid-model' });
const caller = createCaller({ user: { id: 'user-paid', is_admin: false } as User });

await caller.sendMessage({
cloudAgentSessionId: 'agent_123',
payload: { type: 'command', command: 'review', arguments: '' },
});

expect(mockGetSession).toHaveBeenCalledWith('agent_123');
expect(mockComputeCloudAgentNextBalanceCheckEligibility).toHaveBeenCalledWith(
expect.objectContaining({ modelId: 'kilo/paid-model' })
);
expect(mockCreateCloudAgentNextClientForModel).toHaveBeenCalledWith('cloud-agent-token', {
isFree: false,
hasUserByokAvailable: false,
});
});

it('falls back to the balance-checked client when the session model is unavailable', async () => {
mockGetSession.mockResolvedValueOnce({ model: undefined });
const caller = createCaller({ user: { id: 'user-1', is_admin: false } as User });

await caller.sendMessage({
cloudAgentSessionId: 'agent_123',
payload: { type: 'command', command: 'review', arguments: '' },
});

expect(mockGetSession).toHaveBeenCalledWith('agent_123');
expect(mockComputeCloudAgentNextBalanceCheckEligibility).not.toHaveBeenCalled();
expect(mockCreateCloudAgentNextClientForModel).toHaveBeenCalledWith('cloud-agent-token', {
isFree: false,
hasUserByokAvailable: false,
});
});

it('falls back to the balance-checked client when getSession rejects', async () => {
mockGetSession.mockRejectedValueOnce(new Error('worker unavailable'));
const caller = createCaller({ user: { id: 'user-1', is_admin: false } as User });

await caller.sendMessage({
cloudAgentSessionId: 'agent_123',
payload: { type: 'command', command: 'review', arguments: '' },
});

expect(mockGetSession).toHaveBeenCalledWith('agent_123');
expect(mockComputeCloudAgentNextBalanceCheckEligibility).not.toHaveBeenCalled();
expect(mockCreateCloudAgentNextClientForModel).toHaveBeenCalledWith('cloud-agent-token', {
isFree: false,
hasUserByokAvailable: false,
});
expect(mockSendMessage).toHaveBeenCalled();
});

it('signs Cloud Agent document uploads with the authenticated user scope', async () => {
const caller = createCaller({ user: { id: 'user-1', is_admin: false } as User });
await caller.getAttachmentUploadUrl({
Expand Down
29 changes: 28 additions & 1 deletion apps/web/src/routers/cloud-agent-next-router.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -217,7 +217,34 @@ export const cloudAgentNextRouter = createTRPCRouter({
.mutation(async ({ ctx, input }) => {
await assertUserOwnsSession(ctx.user.id, input.cloudAgentSessionId);
const authToken = generateCloudAgentToken(ctx.user);
const client = createCloudAgentNextClient(authToken);
// Prompt turns carry their own model; command turns run the session's
// stored model, so resolve it to apply the same free/BYOK eligibility
// to every follow-up that queues a model-using turn. If the worker
// can't return the session model, fall back to the balance-checked
// client (the safe default) rather than failing the mutation on a
// read-side error.
let modelId: string | undefined;
if (input.payload.type === 'prompt') {
modelId = input.payload.model;
} else {
try {
modelId = (
await createCloudAgentNextClient(authToken).getSession(input.cloudAgentSessionId)
).model;
} catch {
modelId = undefined;
}
}
const client = createCloudAgentNextClientForModel(
authToken,
modelId
? await computeCloudAgentNextBalanceCheckEligibility({
fromDb: db,
user: ctx.user,
modelId,
})
: { isFree: false, hasUserByokAvailable: false }
);

// Tokens are refreshed inside cloud-agent-next (GitHub App installation
// for GitHub, GIT_TOKEN_SERVICE for managed GitLab).
Expand Down
130 changes: 129 additions & 1 deletion apps/web/src/routers/organizations/organization-cloud-agent-next-router.test.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -54,9 +54,12 @@ const mockGenerateCloudAgentAttachmentUploadUrl = jest.fn<
}) => Promise<{ signedUrl: string; key: string; expiresAt: string }>
>(() => Promise.resolve({ signedUrl: 'signed', key: 'key', expiresAt: 'expires' }));

const mockGetSession = jest.fn<(cloudAgentSessionId: string) => Promise<{ model?: string }>>();

const mockCreateCloudAgentNextClient = jest.fn(() => ({
prepareSession: mockPrepareSession,
sendMessage: mockSendMessage,
getSession: mockGetSession,
}));

const mockCreateCloudAgentNextClientForModel = jest.fn(
Expand Down Expand Up @@ -195,7 +198,9 @@ let createCaller: (ctx: { user: User }) => {
sendMessage: (input: {
organizationId: string;
cloudAgentSessionId: string;
payload: { type: 'prompt'; prompt: string; mode: string; model: string };
payload:
| { type: 'prompt'; prompt: string; mode: string; model: string }
| { type: 'command'; command: string; arguments: string };
attachments?: { path: string; files: string[] };
images?: { path: string; files: string[] };
}) => Promise<unknown>;
Expand Down Expand Up @@ -262,6 +267,11 @@ describe('organizationCloudAgentNextRouter attachment forwarding', () => {
mockVerifyOrgOwnsSessionV2ByCloudAgentId.mockResolvedValue({
kiloSessionId: 'ses_12345678901234567890123456',
});
mockComputeCloudAgentNextBalanceCheckEligibility.mockResolvedValue({
isFree: false,
hasUserByokAvailable: false,
});
mockGetSession.mockResolvedValue({ model: 'kilo/paid-model' });
});

it('denies an inaccessible organization session before calling the Worker', async () => {
Expand Down Expand Up @@ -317,6 +327,124 @@ describe('organizationCloudAgentNextRouter attachment forwarding', () => {
expect(mockSendMessage).not.toHaveBeenCalledWith(expect.objectContaining({ images }));
});

it('routes free follow-up prompt models through the balance-skip client', async () => {
mockComputeCloudAgentNextBalanceCheckEligibility.mockResolvedValueOnce({
isFree: true,
hasUserByokAvailable: false,
});
const caller = createCaller({ user: { id: 'user-free', is_admin: false } as User });

await caller.sendMessage({
organizationId: ORGANIZATION_ID,
cloudAgentSessionId: 'agent_123',
payload: {
type: 'prompt',
prompt: 'Follow up on this',
mode: 'code',
model: 'kilo/free-model',
},
});

expect(mockComputeCloudAgentNextBalanceCheckEligibility).toHaveBeenCalledWith(
expect.objectContaining({
modelId: 'kilo/free-model',
organizationId: ORGANIZATION_ID,
})
);
expect(mockCreateCloudAgentNextClientForModel).toHaveBeenCalledWith('cloud-agent-token', {
isFree: true,
hasUserByokAvailable: false,
});
expect(mockCreateCloudAgentNextClient).not.toHaveBeenCalled();
});

it('routes free follow-up command turns through the balance-skip client using the session model', async () => {
mockGetSession.mockResolvedValueOnce({ model: 'kilo/free-model' });
mockComputeCloudAgentNextBalanceCheckEligibility.mockResolvedValueOnce({
isFree: true,
hasUserByokAvailable: false,
});
const caller = createCaller({ user: { id: 'user-free', is_admin: false } as User });

await caller.sendMessage({
organizationId: ORGANIZATION_ID,
cloudAgentSessionId: 'agent_123',
payload: { type: 'command', command: 'review', arguments: '' },
});

expect(mockGetSession).toHaveBeenCalledWith('agent_123');
expect(mockComputeCloudAgentNextBalanceCheckEligibility).toHaveBeenCalledWith(
expect.objectContaining({
modelId: 'kilo/free-model',
organizationId: ORGANIZATION_ID,
})
);
expect(mockCreateCloudAgentNextClientForModel).toHaveBeenCalledWith('cloud-agent-token', {
isFree: true,
hasUserByokAvailable: false,
});
});

it('keeps the balance check for command turns on paid organization sessions', async () => {
mockGetSession.mockResolvedValueOnce({ model: 'kilo/paid-model' });
const caller = createCaller({ user: { id: 'user-paid', is_admin: false } as User });

await caller.sendMessage({
organizationId: ORGANIZATION_ID,
cloudAgentSessionId: 'agent_123',
payload: { type: 'command', command: 'review', arguments: '' },
});

expect(mockGetSession).toHaveBeenCalledWith('agent_123');
expect(mockComputeCloudAgentNextBalanceCheckEligibility).toHaveBeenCalledWith(
expect.objectContaining({
modelId: 'kilo/paid-model',
organizationId: ORGANIZATION_ID,
})
);
expect(mockCreateCloudAgentNextClientForModel).toHaveBeenCalledWith('cloud-agent-token', {
isFree: false,
hasUserByokAvailable: false,
});
});

it('falls back to the balance-checked client when the organization session model is unavailable', async () => {
mockGetSession.mockResolvedValueOnce({ model: undefined });
const caller = createCaller({ user: { id: 'user-1', is_admin: false } as User });

await caller.sendMessage({
organizationId: ORGANIZATION_ID,
cloudAgentSessionId: 'agent_123',
payload: { type: 'command', command: 'review', arguments: '' },
});

expect(mockGetSession).toHaveBeenCalledWith('agent_123');
expect(mockComputeCloudAgentNextBalanceCheckEligibility).not.toHaveBeenCalled();
expect(mockCreateCloudAgentNextClientForModel).toHaveBeenCalledWith('cloud-agent-token', {
isFree: false,
hasUserByokAvailable: false,
});
});

it('falls back to the balance-checked client when getSession rejects', async () => {
mockGetSession.mockRejectedValueOnce(new Error('worker unavailable'));
const caller = createCaller({ user: { id: 'user-1', is_admin: false } as User });

await caller.sendMessage({
organizationId: ORGANIZATION_ID,
cloudAgentSessionId: 'agent_123',
payload: { type: 'command', command: 'review', arguments: '' },
});

expect(mockGetSession).toHaveBeenCalledWith('agent_123');
expect(mockComputeCloudAgentNextBalanceCheckEligibility).not.toHaveBeenCalled();
expect(mockCreateCloudAgentNextClientForModel).toHaveBeenCalledWith('cloud-agent-token', {
isFree: false,
hasUserByokAvailable: false,
});
expect(mockSendMessage).toHaveBeenCalled();
});

it('signs Cloud Agent document uploads within authenticated organization access', async () => {
const caller = createCaller({ user: { id: 'user-1', is_admin: false } as User });
await caller.getAttachmentUploadUrl({
Expand Down
30 changes: 29 additions & 1 deletion apps/web/src/routers/organizations/organization-cloud-agent-next-router.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -331,7 +331,35 @@ export const organizationCloudAgentNextRouter = createTRPCRouter({
cloudAgentSessionId: input.cloudAgentSessionId,
});
const authToken = generateCloudAgentToken(ctx.user);
const client = createCloudAgentNextClient(authToken);
// Prompt turns carry their own model; command turns run the session's
// stored model, so resolve it to apply the same free/BYOK eligibility
// to every follow-up that queues a model-using turn. If the worker
// can't return the session model, fall back to the balance-checked
// client (the safe default) rather than failing the mutation on a
// read-side error.
let modelId: string | undefined;
if (input.payload.type === 'prompt') {
modelId = input.payload.model;
} else {
try {
modelId = (
await createCloudAgentNextClient(authToken).getSession(input.cloudAgentSessionId)
).model;
} catch {
modelId = undefined;
}
}
const client = createCloudAgentNextClientForModel(
authToken,
modelId
? await computeCloudAgentNextBalanceCheckEligibility({
fromDb: db,
user: ctx.user,
modelId,
organizationId: input.organizationId,
})
: { isFree: false, hasUserByokAvailable: false }
);

// Tokens are refreshed inside cloud-agent-next (GitHub App installation
// for GitHub, GIT_TOKEN_SERVICE for managed GitLab). organizationId is
Expand Down