Skip to content

Commit 59f34f6

Browse files
feat(cli): carry ABCA solution UA on all bgagent SDK clients
CLI-local ua.ts mirror (same convention as types.ts — CLI can't import from CDK). Component 'cli', stack name from new optional stack_name config field (configure --stack-name), trace = process pid set once at startup. All Cognito/SecretsManager/CloudFormation/DynamoDB client sites migrated. Wire-capture test mirrors the CDK one via a real Cognito client + stub requestHandler; auth.test.ts asserts the constructor receives the customUserAgent. Task 6 of PR #338 plan. Part of #319 Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
1 parent 352fb18 commit 59f34f6

11 files changed

Lines changed: 359 additions & 18 deletions

File tree

cli/src/auth.ts

Lines changed: 3 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -26,6 +26,7 @@ import { loadConfig, loadCredentials, saveCredentials } from './config';
2626
import { debug } from './debug';
2727
import { CliError } from './errors';
2828
import { Credentials } from './types';
29+
import { abcaUserAgent, withAbcaTrace } from './ua';
2930

3031
const TOKEN_REFRESH_BUFFER_MINUTES = 5;
3132
const TOKEN_REFRESH_BUFFER_MS = TOKEN_REFRESH_BUFFER_MINUTES * 60 * 1000;
@@ -34,7 +35,7 @@ const TOKEN_REFRESH_BUFFER_MS = TOKEN_REFRESH_BUFFER_MINUTES * 60 * 1000;
3435
export async function login(username: string, password: string): Promise<void> {
3536
const config = loadConfig();
3637
debug(`Cognito region: ${config.region}, client_id: ${config.client_id}, user_pool_id: ${config.user_pool_id}`);
37-
const client = new CognitoIdentityProviderClient({ region: config.region });
38+
const client = withAbcaTrace(new CognitoIdentityProviderClient({ region: config.region, ...abcaUserAgent() }));
3839

3940
const result = await client.send(new InitiateAuthCommand({
4041
AuthFlow: AuthFlowType.USER_PASSWORD_AUTH,
@@ -100,7 +101,7 @@ function isExpired(creds: Credentials): boolean {
100101

101102
async function refreshToken(creds: Credentials): Promise<void> {
102103
const config = loadConfig();
103-
const client = new CognitoIdentityProviderClient({ region: config.region });
104+
const client = withAbcaTrace(new CognitoIdentityProviderClient({ region: config.region, ...abcaUserAgent() }));
104105

105106
try {
106107
const result = await client.send(new InitiateAuthCommand({

cli/src/bin/bgagent.ts

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -41,6 +41,11 @@ import { makeWatchCommand } from '../commands/watch';
4141
import { makeWebhookCommand } from '../commands/webhook';
4242
import { setVerbose } from '../debug';
4343
import { ApiError, CliError } from '../errors';
44+
import { setAbcaTrace } from '../ua';
45+
46+
// User-Agent solution tracking (#319): the pid correlates all AWS calls
47+
// made by this CLI invocation.
48+
setAbcaTrace(String(process.pid));
4449

4550
const program = new Command();
4651

cli/src/commands/admin.ts

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -27,6 +27,7 @@ import { Command } from 'commander';
2727
import { loadConfig } from '../config';
2828
import { CliError } from '../errors';
2929
import { CliConfig } from '../types';
30+
import { abcaUserAgent, withAbcaTrace } from '../ua';
3031

3132
/**
3233
* Generate a strong temporary password meeting Cognito's default policy:
@@ -149,7 +150,7 @@ export function makeAdminCommand(): Command {
149150

150151
const tempPassword = opts.tempPassword ?? generateTempPassword();
151152

152-
const cognito = new CognitoIdentityProviderClient({ region });
153+
const cognito = withAbcaTrace(new CognitoIdentityProviderClient({ region, ...abcaUserAgent() }));
153154
try {
154155
await cognito.send(new AdminCreateUserCommand({
155156
UserPoolId: config.user_pool_id,

cli/src/commands/configure.ts

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -37,6 +37,7 @@ export function makeConfigureCommand(): Command {
3737
.option('--region <region>', 'AWS region')
3838
.option('--user-pool-id <id>', 'Cognito User Pool ID')
3939
.option('--client-id <id>', 'Cognito App Client ID')
40+
.option('--stack-name <name>', 'Deployed stack name (User-Agent solution tracking)')
4041
.option('--from-bundle <base64>', 'Base64 config bundle from `bgagent admin invite-user`')
4142
.action((opts) => {
4243
// --from-bundle is mutually exclusive with the individual flags. Mixing
@@ -58,6 +59,7 @@ export function makeConfigureCommand(): Command {
5859
...(opts.region !== undefined ? { region: opts.region } : {}),
5960
...(opts.userPoolId !== undefined ? { user_pool_id: opts.userPoolId } : {}),
6061
...(opts.clientId !== undefined ? { client_id: opts.clientId } : {}),
62+
...(opts.stackName !== undefined ? { stack_name: opts.stackName } : {}),
6163
};
6264
}
6365
const merged: Partial<CliConfig> = {

cli/src/commands/github.ts

Lines changed: 3 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -26,6 +26,7 @@ import {
2626
import { Command } from 'commander';
2727
import { loadConfig } from '../config';
2828
import { CliError } from '../errors';
29+
import { abcaUserAgent, withAbcaTrace } from '../ua';
2930

3031
/** Width of the `═` banner rules printed around webhook-info output. */
3132
const BANNER_WIDTH = 72;
@@ -116,7 +117,7 @@ export function makeGithubCommand(): Command {
116117
);
117118
}
118119

119-
const sm = new SecretsManagerClient({ region });
120+
const sm = withAbcaTrace(new SecretsManagerClient({ region, ...abcaUserAgent() }));
120121

121122
// Show whether a secret is already configured so the operator
122123
// doesn't accidentally rotate it without realising. Linear's
@@ -166,7 +167,7 @@ export function makeGithubCommand(): Command {
166167
// ─── Stack-output helper ─────────────────────────────────────────────────────
167168

168169
async function getStackOutput(region: string, stackName: string, outputKey: string): Promise<string | null> {
169-
const cf = new CloudFormationClient({ region });
170+
const cf = withAbcaTrace(new CloudFormationClient({ region, ...abcaUserAgent() }));
170171
try {
171172
const result = await cf.send(new DescribeStacksCommand({ StackName: stackName }));
172173
const stack = result.Stacks?.[0];

cli/src/commands/linear.ts

Lines changed: 12 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -45,6 +45,7 @@ import {
4545
StoredLinearOauthToken,
4646
} from '../linear-oauth';
4747
import { awaitOauthCallback, CALLBACK_URL } from '../oauth-callback-server';
48+
import { abcaUserAgent, withAbcaTrace } from '../ua';
4849

4950
/** Default label that triggers an ABCA task when applied to a Linear issue. */
5051
const DEFAULT_LABEL_FILTER = 'bgagent';
@@ -597,7 +598,7 @@ export function makeLinearCommand(): Command {
597598

598599
// ─── Step 4: Persist token to per-workspace Secrets Manager ───
599600
process.stdout.write(' → Storing OAuth token...');
600-
const sm = new SecretsManagerClient({ region });
601+
const sm = withAbcaTrace(new SecretsManagerClient({ region, ...abcaUserAgent() }));
601602
const now = new Date().toISOString();
602603
const stored: StoredLinearOauthToken = {
603604
access_token: tokenResponse.access_token,
@@ -625,7 +626,7 @@ export function makeLinearCommand(): Command {
625626
console.log(` ✓ (${secretName})`);
626627

627628
// ─── Step 5: Persist registry + user-mapping rows ─────────────
628-
const ddb = DynamoDBDocumentClient.from(new DynamoDBClient({ region }));
629+
const ddb = DynamoDBDocumentClient.from(withAbcaTrace(new DynamoDBClient({ region, ...abcaUserAgent() })));
629630

630631
// Best-effort: fetch team keys so the screenshot processor can
631632
// prefix-route Linear issue lookups (e.g. ABCA-42 → workspace
@@ -829,8 +830,8 @@ export function makeLinearCommand(): Command {
829830
);
830831
}
831832

832-
const sm = new SecretsManagerClient({ region });
833-
const ddb = DynamoDBDocumentClient.from(new DynamoDBClient({ region }));
833+
const sm = withAbcaTrace(new SecretsManagerClient({ region, ...abcaUserAgent() }));
834+
const ddb = DynamoDBDocumentClient.from(withAbcaTrace(new DynamoDBClient({ region, ...abcaUserAgent() })));
834835

835836
// ─── Linear OAuth app credentials ──────────────────────────────
836837
// Always prompt — never accept secrets via flags (shell history
@@ -1099,7 +1100,7 @@ export function makeLinearCommand(): Command {
10991100
const config = loadConfig();
11001101
const region = opts.region || config.region;
11011102

1102-
const sm = new SecretsManagerClient({ region });
1103+
const sm = withAbcaTrace(new SecretsManagerClient({ region, ...abcaUserAgent() }));
11031104
const secretName = linearOauthSecretName(slug);
11041105

11051106
// ─── Read existing bundle ───────────────────────────────────
@@ -1218,8 +1219,8 @@ export function makeLinearCommand(): Command {
12181219
const callerCognitoSub = extractCognitoSub();
12191220

12201221
// ─── Resolve workspace + OAuth secret arn ──────────────────────
1221-
const sm = new SecretsManagerClient({ region });
1222-
const ddb = DynamoDBDocumentClient.from(new DynamoDBClient({ region }));
1222+
const sm = withAbcaTrace(new SecretsManagerClient({ region, ...abcaUserAgent() }));
1223+
const ddb = DynamoDBDocumentClient.from(withAbcaTrace(new DynamoDBClient({ region, ...abcaUserAgent() })));
12231224
const registryScan = await ddb.send(new ScanCommand({
12241225
TableName: workspaceRegistryTable!,
12251226
FilterExpression: 'workspace_slug = :slug AND #status = :active',
@@ -1342,7 +1343,7 @@ export function makeLinearCommand(): Command {
13421343
}
13431344

13441345
const now = new Date().toISOString();
1345-
const ddb = DynamoDBDocumentClient.from(new DynamoDBClient({ region }));
1346+
const ddb = DynamoDBDocumentClient.from(withAbcaTrace(new DynamoDBClient({ region, ...abcaUserAgent() })));
13461347
await ddb.send(new PutCommand({
13471348
TableName: tableName,
13481349
Item: {
@@ -1373,7 +1374,7 @@ export function makeLinearCommand(): Command {
13731374
.action(async (opts) => {
13741375
const config = loadConfig();
13751376
const region = opts.region || config.region;
1376-
const sm = new SecretsManagerClient({ region });
1377+
const sm = withAbcaTrace(new SecretsManagerClient({ region, ...abcaUserAgent() }));
13771378

13781379
// Resolve the set of workspace slugs to query. Either an
13791380
// explicit `--slug` (one workspace) or every Linear workspace
@@ -1908,7 +1909,7 @@ export async function autoLinkTokenOwner(args: {
19081909
return;
19091910
}
19101911

1911-
const ddb = DynamoDBDocumentClient.from(new DynamoDBClient({ region: args.region }));
1912+
const ddb = DynamoDBDocumentClient.from(withAbcaTrace(new DynamoDBClient({ region: args.region, ...abcaUserAgent() })));
19121913
await ddb.send(new PutCommand({
19131914
TableName: args.userMappingTable,
19141915
Item: {
@@ -1947,7 +1948,7 @@ function extractCognitoSub(): string {
19471948

19481949
async function getStackOutput(region: string, stackName: string, outputKey: string): Promise<string | null> {
19491950
try {
1950-
const cfn = new CloudFormationClient({ region });
1951+
const cfn = withAbcaTrace(new CloudFormationClient({ region, ...abcaUserAgent() }));
19511952
const result = await cfn.send(new DescribeStacksCommand({ StackName: stackName }));
19521953
const outputs = result.Stacks?.[0]?.Outputs ?? [];
19531954
const output = outputs.find((o) => o.OutputKey === outputKey);

cli/src/commands/slack.ts

Lines changed: 3 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -27,6 +27,7 @@ import { Command } from 'commander';
2727
import { ApiClient } from '../api-client';
2828
import { loadConfig } from '../config';
2929
import { formatJson } from '../format';
30+
import { abcaUserAgent, withAbcaTrace } from '../ua';
3031

3132
export function makeSlackCommand(): Command {
3233
const slack = new Command('slack')
@@ -208,7 +209,7 @@ async function promptAndStoreCredentials(region: string, arns: SecretArns): Prom
208209

209210
// Store in Secrets Manager.
210211
console.log('');
211-
const sm = new SecretsManagerClient({ region });
212+
const sm = withAbcaTrace(new SecretsManagerClient({ region, ...abcaUserAgent() }));
212213

213214
const secrets = [
214215
{ id: arns.signingSecretArn, value: signingSecret, label: 'signing secret' },
@@ -345,7 +346,7 @@ function findRepoRoot(): string {
345346

346347
async function getStackOutput(region: string, stackName: string, outputKey: string): Promise<string | null> {
347348
try {
348-
const cfn = new CloudFormationClient({ region });
349+
const cfn = withAbcaTrace(new CloudFormationClient({ region, ...abcaUserAgent() }));
349350
const result = await cfn.send(new DescribeStacksCommand({ StackName: stackName }));
350351
const outputs = result.Stacks?.[0]?.Outputs ?? [];
351352
const output = outputs.find((o) => o.OutputKey === outputKey);

cli/src/types.ts

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -372,6 +372,8 @@ export interface CliConfig {
372372
readonly region: string;
373373
readonly user_pool_id: string;
374374
readonly client_id: string;
375+
/** Deployed stack name for User-Agent solution tracking (#319); optional. */
376+
readonly stack_name?: string;
375377
}
376378

377379
/** Cached credentials stored in ~/.bgagent/credentials.json.

cli/src/ua.ts

Lines changed: 137 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,137 @@
1+
/**
2+
* MIT No Attribution
3+
*
4+
* Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
5+
*
6+
* Permission is hereby granted, free of charge, to any person obtaining a copy of
7+
* the Software without restriction, including without limitation the rights to
8+
* use, copy, modify, merge, publish, distribute, sublicense, and/or sell copies of
9+
* the Software, and to permit persons to whom the Software is furnished to do so.
10+
*
11+
* THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
12+
* IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
13+
* FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
14+
* AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
15+
* LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
16+
* OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
17+
* SOFTWARE.
18+
*/
19+
20+
/**
21+
* Outbound AWS SDK User-Agent solution tracking (#319) — CLI surface.
22+
*
23+
* Every AWS API call made by `bgagent` carries:
24+
*
25+
* app/uksb-wt64nei4u6/{STACKNAME} (only when config has stack_name)
26+
* md/uksb-wt64nei4u6#cli[#{PID}]
27+
*
28+
* CLI-local mirror of `cdk/src/handlers/shared/ua.ts` (the CLI package
29+
* cannot import from the CDK package — same mirroring convention as
30+
* `cli/src/types.ts`). Solution id, wire format, and sanitization rules
31+
* must stay identical; `agent/src/ua.py` is the Python counterpart.
32+
*
33+
* The component label is hardcoded (`cli`); the stack name comes from the
34+
* optional `stack_name` field in `~/.bgagent/config.json`. The trace handle
35+
* is the CLI process pid, set once at startup in `bin/bgagent.ts` and
36+
* appended per-request by the {@link withAbcaTrace} middleware.
37+
*/
38+
39+
import { tryLoadConfig } from './config';
40+
41+
/**
42+
* AWS solution-tracking id for ABCA. Deploy-time counterpart (#292) lives in
43+
* the CloudFormation stack description in `cdk/src/main.ts`.
44+
*/
45+
export const SOLUTION_ID = 'uksb-wt64nei4u6';
46+
47+
/** Stable per-component label: this surface IS the bgagent CLI. */
48+
const COMPONENT = 'cli';
49+
50+
/** App-id budget: 50-char value cap minus `uksb-wt64nei4u6/` (16) = 34. */
51+
const STACK_NAME_MAX = 34;
52+
53+
/**
54+
* RFC 7230 token charset; `/` and `#` deliberately excluded (structural
55+
* separators of the scheme). Mirrors the CDK and Python implementations.
56+
*/
57+
const UA_TOKEN_SAFE = /[^A-Za-z0-9!$%&'*+\-.^_`|~]/g;
58+
59+
let currentTrace: string | undefined;
60+
61+
/** Replace every non-UA-token char (incl. non-ASCII) with `-`. */
62+
export function sanitizeUaValue(raw: string): string {
63+
return raw.replace(UA_TOKEN_SAFE, '-');
64+
}
65+
66+
/**
67+
* Client config fragment carrying the static ABCA UA segments. Spread into
68+
* every SDK client constructor: `new SecretsManagerClient({ region, ...abcaUserAgent() })`.
69+
*
70+
* Pair semantics (mirrors the CDK module): the `app/` segment is a
71+
* single-element pair so its literal `/` separators survive the SDK's
72+
* name-position escaping; the `md/` pair lets the SDK's own `name#value`
73+
* join produce the `#`.
74+
*/
75+
export function abcaUserAgent(): { customUserAgent: ([string] | [string, string])[] } {
76+
const pairs: ([string] | [string, string])[] = [];
77+
const stackName = tryLoadConfig()?.stack_name?.trim();
78+
if (stackName) {
79+
// Sanitize FIRST, then clip, so a replaced char can't be re-split.
80+
const clipped = sanitizeUaValue(stackName).slice(0, STACK_NAME_MAX);
81+
pairs.push([`app/${SOLUTION_ID}/${clipped}`]);
82+
}
83+
pairs.push([`md/${SOLUTION_ID}`, COMPONENT]);
84+
return { customUserAgent: pairs };
85+
}
86+
87+
/** Set (or clear) the ambient trace handle (the CLI pid, set at startup). */
88+
export function setAbcaTrace(handle?: string): void {
89+
currentTrace = handle || undefined;
90+
}
91+
92+
/** Current trace handle, sanitized to UA-token-safe ASCII, or undefined. */
93+
export function getAbcaTrace(): string | undefined {
94+
return currentTrace ? sanitizeUaValue(currentTrace) : undefined;
95+
}
96+
97+
/** Structural view of a client middleware stack (avoids @smithy/types dep). */
98+
interface MiddlewareStackLike {
99+
addRelativeTo(middleware: unknown, options: Record<string, unknown>): void;
100+
}
101+
102+
/**
103+
* Append `#{TRACE}` to the outgoing User-Agent headers on every request by
104+
* splicing onto the static `md/` segment, after the SDK's own
105+
* `getUserAgentMiddleware` has rendered the headers. Mutates only the
106+
* header strings; the client and its connection pool are untouched.
107+
* No-ops on clients without a middleware stack (jest constructor mocks).
108+
*/
109+
export function withAbcaTrace<T>(client: T): T {
110+
const stack = (client as { middlewareStack?: MiddlewareStackLike }).middlewareStack;
111+
if (!stack || typeof stack.addRelativeTo !== 'function') {
112+
return client;
113+
}
114+
const md = `md/${SOLUTION_ID}#${COMPONENT}`;
115+
stack.addRelativeTo(
116+
(next: (args: unknown) => Promise<unknown>) => async (args: unknown) => {
117+
const trace = getAbcaTrace();
118+
const request = (args as { request?: { headers?: Record<string, string> } }).request;
119+
if (trace && request?.headers) {
120+
for (const header of ['user-agent', 'x-amz-user-agent']) {
121+
const value = request.headers[header];
122+
if (value && value.includes(md)) {
123+
request.headers[header] = value.replace(md, `${md}#${trace}`);
124+
}
125+
}
126+
}
127+
return next(args);
128+
},
129+
{
130+
name: 'abcaUaTraceMiddleware',
131+
relation: 'after',
132+
toMiddleware: 'getUserAgentMiddleware',
133+
override: true,
134+
},
135+
);
136+
return client;
137+
}

cli/test/auth.test.ts

Lines changed: 13 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -52,6 +52,19 @@ describe('auth', () => {
5252
});
5353

5454
describe('login', () => {
55+
test('Cognito client carries the ABCA solution customUserAgent (#319)', async () => {
56+
const { CognitoIdentityProviderClient } = jest.requireMock(
57+
'@aws-sdk/client-cognito-identity-provider',
58+
);
59+
mockSend.mockResolvedValue({
60+
AuthenticationResult: { IdToken: 'i', AccessToken: 'a', RefreshToken: 'r', ExpiresIn: 3600 },
61+
});
62+
await login('user', 'pass');
63+
const calls = (CognitoIdentityProviderClient as jest.Mock).mock.calls;
64+
const config = calls[calls.length - 1]?.[0];
65+
expect(config.customUserAgent).toEqual([['md/uksb-wt64nei4u6', 'cli']]);
66+
});
67+
5568
test('saves credentials on successful login', async () => {
5669
mockSend.mockResolvedValue({
5770
AuthenticationResult: {

0 commit comments

Comments
 (0)