diff --git a/codex-rs/app-server-protocol/schema/json/ClientRequest.json b/codex-rs/app-server-protocol/schema/json/ClientRequest.json index 97b5f698c2..eace6abe71 100644 --- a/codex-rs/app-server-protocol/schema/json/ClientRequest.json +++ b/codex-rs/app-server-protocol/schema/json/ClientRequest.json @@ -4030,6 +4030,24 @@ ], "type": "object" }, + "UsageRange": { + "enum": [ + "day", + "week" + ], + "type": "string" + }, + "UsageReadParams": { + "properties": { + "range": { + "$ref": "#/definitions/UsageRange" + } + }, + "required": [ + "range" + ], + "type": "object" + }, "UserInput": { "oneOf": [ { @@ -4829,6 +4847,30 @@ "title": "Plugin/listRequest", "type": "object" }, + { + "properties": { + "id": { + "$ref": "#/definitions/RequestId" + }, + "method": { + "enum": [ + "usage/read" + ], + "title": "Usage/readRequestMethod", + "type": "string" + }, + "params": { + "$ref": "#/definitions/UsageReadParams" + } + }, + "required": [ + "id", + "method", + "params" + ], + "title": "Usage/readRequest", + "type": "object" + }, { "properties": { "id": { diff --git a/codex-rs/app-server-protocol/schema/json/codex_app_server_protocol.schemas.json b/codex-rs/app-server-protocol/schema/json/codex_app_server_protocol.schemas.json index dda285476d..402e8932f9 100644 --- a/codex-rs/app-server-protocol/schema/json/codex_app_server_protocol.schemas.json +++ b/codex-rs/app-server-protocol/schema/json/codex_app_server_protocol.schemas.json @@ -829,6 +829,30 @@ "title": "Plugin/listRequest", "type": "object" }, + { + "properties": { + "id": { + "$ref": "#/definitions/v2/RequestId" + }, + "method": { + "enum": [ + "usage/read" + ], + "title": "Usage/readRequestMethod", + "type": "string" + }, + "params": { + "$ref": "#/definitions/v2/UsageReadParams" + } + }, + "required": [ + "id", + "method", + "params" + ], + "title": "Usage/readRequest", + "type": "object" + }, { "properties": { "id": { @@ -18555,6 +18579,169 @@ "title": "TurnSteerResponse", "type": "object" }, + "UsageContributorKind": { + "enum": [ + "skill", + "subagent", + "app", + "mcpServer", + "plugin" + ], + "type": "string" + }, + "UsageEntry": { + "properties": { + "attributedTokens": { + "format": "int64", + "type": "integer" + }, + "id": { + "type": "string" + }, + "kind": { + "$ref": "#/definitions/v2/UsageContributorKind" + }, + "label": { + "type": "string" + }, + "percentOfUsage": { + "format": "uint8", + "minimum": 0.0, + "type": "integer" + } + }, + "required": [ + "attributedTokens", + "id", + "kind", + "label", + "percentOfUsage" + ], + "type": "object" + }, + "UsageHeadline": { + "properties": { + "entry": { + "$ref": "#/definitions/v2/UsageEntry" + }, + "note": { + "type": [ + "string", + "null" + ] + } + }, + "required": [ + "entry" + ], + "type": "object" + }, + "UsageRange": { + "enum": [ + "day", + "week" + ], + "type": "string" + }, + "UsageReadParams": { + "$schema": "http://json-schema.org/draft-07/schema#", + "properties": { + "range": { + "$ref": "#/definitions/v2/UsageRange" + } + }, + "required": [ + "range" + ], + "title": "UsageReadParams", + "type": "object" + }, + "UsageReadResponse": { + "$schema": "http://json-schema.org/draft-07/schema#", + "properties": { + "report": { + "$ref": "#/definitions/v2/UsageReport" + } + }, + "required": [ + "report" + ], + "title": "UsageReadResponse", + "type": "object" + }, + "UsageReport": { + "properties": { + "apps": { + "items": { + "$ref": "#/definitions/v2/UsageEntry" + }, + "type": "array" + }, + "generatedAt": { + "format": "int64", + "type": "integer" + }, + "headline": { + "anyOf": [ + { + "$ref": "#/definitions/v2/UsageHeadline" + }, + { + "type": "null" + } + ] + }, + "mcpServers": { + "items": { + "$ref": "#/definitions/v2/UsageEntry" + }, + "type": "array" + }, + "plugins": { + "items": { + "$ref": "#/definitions/v2/UsageEntry" + }, + "type": "array" + }, + "range": { + "$ref": "#/definitions/v2/UsageRange" + }, + "skills": { + "items": { + "$ref": "#/definitions/v2/UsageEntry" + }, + "type": "array" + }, + "subagents": { + "items": { + "$ref": "#/definitions/v2/UsageEntry" + }, + "type": "array" + }, + "totalTokens": { + "format": "int64", + "type": "integer" + }, + "trackedFrom": { + "format": "int64", + "type": [ + "integer", + "null" + ] + } + }, + "required": [ + "apps", + "generatedAt", + "mcpServers", + "plugins", + "range", + "skills", + "subagents", + "totalTokens" + ], + "type": "object" + }, "UserInput": { "oneOf": [ { diff --git a/codex-rs/app-server-protocol/schema/json/codex_app_server_protocol.v2.schemas.json b/codex-rs/app-server-protocol/schema/json/codex_app_server_protocol.v2.schemas.json index 3a23ba3a69..406f7a6fee 100644 --- a/codex-rs/app-server-protocol/schema/json/codex_app_server_protocol.v2.schemas.json +++ b/codex-rs/app-server-protocol/schema/json/codex_app_server_protocol.v2.schemas.json @@ -1555,6 +1555,30 @@ "title": "Plugin/listRequest", "type": "object" }, + { + "properties": { + "id": { + "$ref": "#/definitions/RequestId" + }, + "method": { + "enum": [ + "usage/read" + ], + "title": "Usage/readRequestMethod", + "type": "string" + }, + "params": { + "$ref": "#/definitions/UsageReadParams" + } + }, + "required": [ + "id", + "method", + "params" + ], + "title": "Usage/readRequest", + "type": "object" + }, { "properties": { "id": { @@ -16379,6 +16403,169 @@ "title": "TurnSteerResponse", "type": "object" }, + "UsageContributorKind": { + "enum": [ + "skill", + "subagent", + "app", + "mcpServer", + "plugin" + ], + "type": "string" + }, + "UsageEntry": { + "properties": { + "attributedTokens": { + "format": "int64", + "type": "integer" + }, + "id": { + "type": "string" + }, + "kind": { + "$ref": "#/definitions/UsageContributorKind" + }, + "label": { + "type": "string" + }, + "percentOfUsage": { + "format": "uint8", + "minimum": 0.0, + "type": "integer" + } + }, + "required": [ + "attributedTokens", + "id", + "kind", + "label", + "percentOfUsage" + ], + "type": "object" + }, + "UsageHeadline": { + "properties": { + "entry": { + "$ref": "#/definitions/UsageEntry" + }, + "note": { + "type": [ + "string", + "null" + ] + } + }, + "required": [ + "entry" + ], + "type": "object" + }, + "UsageRange": { + "enum": [ + "day", + "week" + ], + "type": "string" + }, + "UsageReadParams": { + "$schema": "http://json-schema.org/draft-07/schema#", + "properties": { + "range": { + "$ref": "#/definitions/UsageRange" + } + }, + "required": [ + "range" + ], + "title": "UsageReadParams", + "type": "object" + }, + "UsageReadResponse": { + "$schema": "http://json-schema.org/draft-07/schema#", + "properties": { + "report": { + "$ref": "#/definitions/UsageReport" + } + }, + "required": [ + "report" + ], + "title": "UsageReadResponse", + "type": "object" + }, + "UsageReport": { + "properties": { + "apps": { + "items": { + "$ref": "#/definitions/UsageEntry" + }, + "type": "array" + }, + "generatedAt": { + "format": "int64", + "type": "integer" + }, + "headline": { + "anyOf": [ + { + "$ref": "#/definitions/UsageHeadline" + }, + { + "type": "null" + } + ] + }, + "mcpServers": { + "items": { + "$ref": "#/definitions/UsageEntry" + }, + "type": "array" + }, + "plugins": { + "items": { + "$ref": "#/definitions/UsageEntry" + }, + "type": "array" + }, + "range": { + "$ref": "#/definitions/UsageRange" + }, + "skills": { + "items": { + "$ref": "#/definitions/UsageEntry" + }, + "type": "array" + }, + "subagents": { + "items": { + "$ref": "#/definitions/UsageEntry" + }, + "type": "array" + }, + "totalTokens": { + "format": "int64", + "type": "integer" + }, + "trackedFrom": { + "format": "int64", + "type": [ + "integer", + "null" + ] + } + }, + "required": [ + "apps", + "generatedAt", + "mcpServers", + "plugins", + "range", + "skills", + "subagents", + "totalTokens" + ], + "type": "object" + }, "UserInput": { "oneOf": [ { diff --git a/codex-rs/app-server-protocol/schema/json/v2/UsageReadParams.json b/codex-rs/app-server-protocol/schema/json/v2/UsageReadParams.json new file mode 100644 index 0000000000..a0833295d7 --- /dev/null +++ b/codex-rs/app-server-protocol/schema/json/v2/UsageReadParams.json @@ -0,0 +1,22 @@ +{ + "$schema": "http://json-schema.org/draft-07/schema#", + "definitions": { + "UsageRange": { + "enum": [ + "day", + "week" + ], + "type": "string" + } + }, + "properties": { + "range": { + "$ref": "#/definitions/UsageRange" + } + }, + "required": [ + "range" + ], + "title": "UsageReadParams", + "type": "object" +} \ No newline at end of file diff --git a/codex-rs/app-server-protocol/schema/json/v2/UsageReadResponse.json b/codex-rs/app-server-protocol/schema/json/v2/UsageReadResponse.json new file mode 100644 index 0000000000..4f74b47896 --- /dev/null +++ b/codex-rs/app-server-protocol/schema/json/v2/UsageReadResponse.json @@ -0,0 +1,152 @@ +{ + "$schema": "http://json-schema.org/draft-07/schema#", + "definitions": { + "UsageContributorKind": { + "enum": [ + "skill", + "subagent", + "app", + "mcpServer", + "plugin" + ], + "type": "string" + }, + "UsageEntry": { + "properties": { + "attributedTokens": { + "format": "int64", + "type": "integer" + }, + "id": { + "type": "string" + }, + "kind": { + "$ref": "#/definitions/UsageContributorKind" + }, + "label": { + "type": "string" + }, + "percentOfUsage": { + "format": "uint8", + "minimum": 0.0, + "type": "integer" + } + }, + "required": [ + "attributedTokens", + "id", + "kind", + "label", + "percentOfUsage" + ], + "type": "object" + }, + "UsageHeadline": { + "properties": { + "entry": { + "$ref": "#/definitions/UsageEntry" + }, + "note": { + "type": [ + "string", + "null" + ] + } + }, + "required": [ + "entry" + ], + "type": "object" + }, + "UsageRange": { + "enum": [ + "day", + "week" + ], + "type": "string" + }, + "UsageReport": { + "properties": { + "apps": { + "items": { + "$ref": "#/definitions/UsageEntry" + }, + "type": "array" + }, + "generatedAt": { + "format": "int64", + "type": "integer" + }, + "headline": { + "anyOf": [ + { + "$ref": "#/definitions/UsageHeadline" + }, + { + "type": "null" + } + ] + }, + "mcpServers": { + "items": { + "$ref": "#/definitions/UsageEntry" + }, + "type": "array" + }, + "plugins": { + "items": { + "$ref": "#/definitions/UsageEntry" + }, + "type": "array" + }, + "range": { + "$ref": "#/definitions/UsageRange" + }, + "skills": { + "items": { + "$ref": "#/definitions/UsageEntry" + }, + "type": "array" + }, + "subagents": { + "items": { + "$ref": "#/definitions/UsageEntry" + }, + "type": "array" + }, + "totalTokens": { + "format": "int64", + "type": "integer" + }, + "trackedFrom": { + "format": "int64", + "type": [ + "integer", + "null" + ] + } + }, + "required": [ + "apps", + "generatedAt", + "mcpServers", + "plugins", + "range", + "skills", + "subagents", + "totalTokens" + ], + "type": "object" + } + }, + "properties": { + "report": { + "$ref": "#/definitions/UsageReport" + } + }, + "required": [ + "report" + ], + "title": "UsageReadResponse", + "type": "object" +} \ No newline at end of file diff --git a/codex-rs/app-server-protocol/schema/typescript/ClientRequest.ts b/codex-rs/app-server-protocol/schema/typescript/ClientRequest.ts index 86ba2cd171..03b070af24 100644 --- a/codex-rs/app-server-protocol/schema/typescript/ClientRequest.ts +++ b/codex-rs/app-server-protocol/schema/typescript/ClientRequest.ts @@ -80,9 +80,10 @@ import type { ThreadUnsubscribeParams } from "./v2/ThreadUnsubscribeParams"; import type { TurnInterruptParams } from "./v2/TurnInterruptParams"; import type { TurnStartParams } from "./v2/TurnStartParams"; import type { TurnSteerParams } from "./v2/TurnSteerParams"; +import type { UsageReadParams } from "./v2/UsageReadParams"; import type { WindowsSandboxSetupStartParams } from "./v2/WindowsSandboxSetupStartParams"; /** * Request from the client to the server. */ -export type ClientRequest ={ "method": "initialize", id: RequestId, params: InitializeParams, } | { "method": "thread/start", id: RequestId, params: ThreadStartParams, } | { "method": "thread/resume", id: RequestId, params: ThreadResumeParams, } | { "method": "thread/fork", id: RequestId, params: ThreadForkParams, } | { "method": "thread/archive", id: RequestId, params: ThreadArchiveParams, } | { "method": "thread/unsubscribe", id: RequestId, params: ThreadUnsubscribeParams, } | { "method": "thread/name/set", id: RequestId, params: ThreadSetNameParams, } | { "method": "thread/goal/set", id: RequestId, params: ThreadGoalSetParams, } | { "method": "thread/goal/get", id: RequestId, params: ThreadGoalGetParams, } | { "method": "thread/goal/clear", id: RequestId, params: ThreadGoalClearParams, } | { "method": "thread/metadata/update", id: RequestId, params: ThreadMetadataUpdateParams, } | { "method": "thread/unarchive", id: RequestId, params: ThreadUnarchiveParams, } | { "method": "thread/compact/start", id: RequestId, params: ThreadCompactStartParams, } | { "method": "thread/shellCommand", id: RequestId, params: ThreadShellCommandParams, } | { "method": "thread/approveGuardianDeniedAction", id: RequestId, params: ThreadApproveGuardianDeniedActionParams, } | { "method": "thread/rollback", id: RequestId, params: ThreadRollbackParams, } | { "method": "thread/list", id: RequestId, params: ThreadListParams, } | { "method": "thread/loaded/list", id: RequestId, params: ThreadLoadedListParams, } | { "method": "thread/read", id: RequestId, params: ThreadReadParams, } | { "method": "thread/inject_items", id: RequestId, params: ThreadInjectItemsParams, } | { "method": "skills/list", id: RequestId, params: SkillsListParams, } | { "method": "hooks/list", id: RequestId, params: HooksListParams, } | { "method": "marketplace/add", id: RequestId, params: MarketplaceAddParams, } | { "method": "marketplace/remove", id: RequestId, params: MarketplaceRemoveParams, } | { "method": "marketplace/upgrade", id: RequestId, params: MarketplaceUpgradeParams, } | { "method": "plugin/list", id: RequestId, params: PluginListParams, } | { "method": "plugin/installed", id: RequestId, params: PluginInstalledParams, } | { "method": "plugin/read", id: RequestId, params: PluginReadParams, } | { "method": "plugin/skill/read", id: RequestId, params: PluginSkillReadParams, } | { "method": "plugin/share/save", id: RequestId, params: PluginShareSaveParams, } | { "method": "plugin/share/updateTargets", id: RequestId, params: PluginShareUpdateTargetsParams, } | { "method": "plugin/share/list", id: RequestId, params: PluginShareListParams, } | { "method": "plugin/share/checkout", id: RequestId, params: PluginShareCheckoutParams, } | { "method": "plugin/share/delete", id: RequestId, params: PluginShareDeleteParams, } | { "method": "app/list", id: RequestId, params: AppsListParams, } | { "method": "fs/readFile", id: RequestId, params: FsReadFileParams, } | { "method": "fs/writeFile", id: RequestId, params: FsWriteFileParams, } | { "method": "fs/createDirectory", id: RequestId, params: FsCreateDirectoryParams, } | { "method": "fs/getMetadata", id: RequestId, params: FsGetMetadataParams, } | { "method": "fs/readDirectory", id: RequestId, params: FsReadDirectoryParams, } | { "method": "fs/remove", id: RequestId, params: FsRemoveParams, } | { "method": "fs/copy", id: RequestId, params: FsCopyParams, } | { "method": "fs/watch", id: RequestId, params: FsWatchParams, } | { "method": "fs/unwatch", id: RequestId, params: FsUnwatchParams, } | { "method": "skills/config/write", id: RequestId, params: SkillsConfigWriteParams, } | { "method": "plugin/install", id: RequestId, params: PluginInstallParams, } | { "method": "plugin/uninstall", id: RequestId, params: PluginUninstallParams, } | { "method": "turn/start", id: RequestId, params: TurnStartParams, } | { "method": "turn/steer", id: RequestId, params: TurnSteerParams, } | { "method": "turn/interrupt", id: RequestId, params: TurnInterruptParams, } | { "method": "review/start", id: RequestId, params: ReviewStartParams, } | { "method": "model/list", id: RequestId, params: ModelListParams, } | { "method": "modelProvider/capabilities/read", id: RequestId, params: ModelProviderCapabilitiesReadParams, } | { "method": "experimentalFeature/list", id: RequestId, params: ExperimentalFeatureListParams, } | { "method": "permissionProfile/list", id: RequestId, params: PermissionProfileListParams, } | { "method": "experimentalFeature/enablement/set", id: RequestId, params: ExperimentalFeatureEnablementSetParams, } | { "method": "mcpServer/oauth/login", id: RequestId, params: McpServerOauthLoginParams, } | { "method": "config/mcpServer/reload", id: RequestId, params: undefined, } | { "method": "mcpServerStatus/list", id: RequestId, params: ListMcpServerStatusParams, } | { "method": "mcpServer/resource/read", id: RequestId, params: McpResourceReadParams, } | { "method": "mcpServer/tool/call", id: RequestId, params: McpServerToolCallParams, } | { "method": "windowsSandbox/setupStart", id: RequestId, params: WindowsSandboxSetupStartParams, } | { "method": "windowsSandbox/readiness", id: RequestId, params: undefined, } | { "method": "account/login/start", id: RequestId, params: LoginAccountParams, } | { "method": "account/login/cancel", id: RequestId, params: CancelLoginAccountParams, } | { "method": "account/logout", id: RequestId, params: undefined, } | { "method": "account/rateLimits/read", id: RequestId, params: undefined, } | { "method": "account/sendAddCreditsNudgeEmail", id: RequestId, params: SendAddCreditsNudgeEmailParams, } | { "method": "feedback/upload", id: RequestId, params: FeedbackUploadParams, } | { "method": "command/exec", id: RequestId, params: CommandExecParams, } | { "method": "command/exec/write", id: RequestId, params: CommandExecWriteParams, } | { "method": "command/exec/terminate", id: RequestId, params: CommandExecTerminateParams, } | { "method": "command/exec/resize", id: RequestId, params: CommandExecResizeParams, } | { "method": "config/read", id: RequestId, params: ConfigReadParams, } | { "method": "externalAgentConfig/detect", id: RequestId, params: ExternalAgentConfigDetectParams, } | { "method": "externalAgentConfig/import", id: RequestId, params: ExternalAgentConfigImportParams, } | { "method": "config/value/write", id: RequestId, params: ConfigValueWriteParams, } | { "method": "config/batchWrite", id: RequestId, params: ConfigBatchWriteParams, } | { "method": "configRequirements/read", id: RequestId, params: undefined, } | { "method": "account/read", id: RequestId, params: GetAccountParams, } | { "method": "getConversationSummary", id: RequestId, params: GetConversationSummaryParams, } | { "method": "gitDiffToRemote", id: RequestId, params: GitDiffToRemoteParams, } | { "method": "getAuthStatus", id: RequestId, params: GetAuthStatusParams, } | { "method": "fuzzyFileSearch", id: RequestId, params: FuzzyFileSearchParams, }; +export type ClientRequest ={ "method": "initialize", id: RequestId, params: InitializeParams, } | { "method": "thread/start", id: RequestId, params: ThreadStartParams, } | { "method": "thread/resume", id: RequestId, params: ThreadResumeParams, } | { "method": "thread/fork", id: RequestId, params: ThreadForkParams, } | { "method": "thread/archive", id: RequestId, params: ThreadArchiveParams, } | { "method": "thread/unsubscribe", id: RequestId, params: ThreadUnsubscribeParams, } | { "method": "thread/name/set", id: RequestId, params: ThreadSetNameParams, } | { "method": "thread/goal/set", id: RequestId, params: ThreadGoalSetParams, } | { "method": "thread/goal/get", id: RequestId, params: ThreadGoalGetParams, } | { "method": "thread/goal/clear", id: RequestId, params: ThreadGoalClearParams, } | { "method": "thread/metadata/update", id: RequestId, params: ThreadMetadataUpdateParams, } | { "method": "thread/unarchive", id: RequestId, params: ThreadUnarchiveParams, } | { "method": "thread/compact/start", id: RequestId, params: ThreadCompactStartParams, } | { "method": "thread/shellCommand", id: RequestId, params: ThreadShellCommandParams, } | { "method": "thread/approveGuardianDeniedAction", id: RequestId, params: ThreadApproveGuardianDeniedActionParams, } | { "method": "thread/rollback", id: RequestId, params: ThreadRollbackParams, } | { "method": "thread/list", id: RequestId, params: ThreadListParams, } | { "method": "thread/loaded/list", id: RequestId, params: ThreadLoadedListParams, } | { "method": "thread/read", id: RequestId, params: ThreadReadParams, } | { "method": "thread/inject_items", id: RequestId, params: ThreadInjectItemsParams, } | { "method": "skills/list", id: RequestId, params: SkillsListParams, } | { "method": "hooks/list", id: RequestId, params: HooksListParams, } | { "method": "marketplace/add", id: RequestId, params: MarketplaceAddParams, } | { "method": "marketplace/remove", id: RequestId, params: MarketplaceRemoveParams, } | { "method": "marketplace/upgrade", id: RequestId, params: MarketplaceUpgradeParams, } | { "method": "plugin/list", id: RequestId, params: PluginListParams, } | { "method": "usage/read", id: RequestId, params: UsageReadParams, } | { "method": "plugin/installed", id: RequestId, params: PluginInstalledParams, } | { "method": "plugin/read", id: RequestId, params: PluginReadParams, } | { "method": "plugin/skill/read", id: RequestId, params: PluginSkillReadParams, } | { "method": "plugin/share/save", id: RequestId, params: PluginShareSaveParams, } | { "method": "plugin/share/updateTargets", id: RequestId, params: PluginShareUpdateTargetsParams, } | { "method": "plugin/share/list", id: RequestId, params: PluginShareListParams, } | { "method": "plugin/share/checkout", id: RequestId, params: PluginShareCheckoutParams, } | { "method": "plugin/share/delete", id: RequestId, params: PluginShareDeleteParams, } | { "method": "app/list", id: RequestId, params: AppsListParams, } | { "method": "fs/readFile", id: RequestId, params: FsReadFileParams, } | { "method": "fs/writeFile", id: RequestId, params: FsWriteFileParams, } | { "method": "fs/createDirectory", id: RequestId, params: FsCreateDirectoryParams, } | { "method": "fs/getMetadata", id: RequestId, params: FsGetMetadataParams, } | { "method": "fs/readDirectory", id: RequestId, params: FsReadDirectoryParams, } | { "method": "fs/remove", id: RequestId, params: FsRemoveParams, } | { "method": "fs/copy", id: RequestId, params: FsCopyParams, } | { "method": "fs/watch", id: RequestId, params: FsWatchParams, } | { "method": "fs/unwatch", id: RequestId, params: FsUnwatchParams, } | { "method": "skills/config/write", id: RequestId, params: SkillsConfigWriteParams, } | { "method": "plugin/install", id: RequestId, params: PluginInstallParams, } | { "method": "plugin/uninstall", id: RequestId, params: PluginUninstallParams, } | { "method": "turn/start", id: RequestId, params: TurnStartParams, } | { "method": "turn/steer", id: RequestId, params: TurnSteerParams, } | { "method": "turn/interrupt", id: RequestId, params: TurnInterruptParams, } | { "method": "review/start", id: RequestId, params: ReviewStartParams, } | { "method": "model/list", id: RequestId, params: ModelListParams, } | { "method": "modelProvider/capabilities/read", id: RequestId, params: ModelProviderCapabilitiesReadParams, } | { "method": "experimentalFeature/list", id: RequestId, params: ExperimentalFeatureListParams, } | { "method": "permissionProfile/list", id: RequestId, params: PermissionProfileListParams, } | { "method": "experimentalFeature/enablement/set", id: RequestId, params: ExperimentalFeatureEnablementSetParams, } | { "method": "mcpServer/oauth/login", id: RequestId, params: McpServerOauthLoginParams, } | { "method": "config/mcpServer/reload", id: RequestId, params: undefined, } | { "method": "mcpServerStatus/list", id: RequestId, params: ListMcpServerStatusParams, } | { "method": "mcpServer/resource/read", id: RequestId, params: McpResourceReadParams, } | { "method": "mcpServer/tool/call", id: RequestId, params: McpServerToolCallParams, } | { "method": "windowsSandbox/setupStart", id: RequestId, params: WindowsSandboxSetupStartParams, } | { "method": "windowsSandbox/readiness", id: RequestId, params: undefined, } | { "method": "account/login/start", id: RequestId, params: LoginAccountParams, } | { "method": "account/login/cancel", id: RequestId, params: CancelLoginAccountParams, } | { "method": "account/logout", id: RequestId, params: undefined, } | { "method": "account/rateLimits/read", id: RequestId, params: undefined, } | { "method": "account/sendAddCreditsNudgeEmail", id: RequestId, params: SendAddCreditsNudgeEmailParams, } | { "method": "feedback/upload", id: RequestId, params: FeedbackUploadParams, } | { "method": "command/exec", id: RequestId, params: CommandExecParams, } | { "method": "command/exec/write", id: RequestId, params: CommandExecWriteParams, } | { "method": "command/exec/terminate", id: RequestId, params: CommandExecTerminateParams, } | { "method": "command/exec/resize", id: RequestId, params: CommandExecResizeParams, } | { "method": "config/read", id: RequestId, params: ConfigReadParams, } | { "method": "externalAgentConfig/detect", id: RequestId, params: ExternalAgentConfigDetectParams, } | { "method": "externalAgentConfig/import", id: RequestId, params: ExternalAgentConfigImportParams, } | { "method": "config/value/write", id: RequestId, params: ConfigValueWriteParams, } | { "method": "config/batchWrite", id: RequestId, params: ConfigBatchWriteParams, } | { "method": "configRequirements/read", id: RequestId, params: undefined, } | { "method": "account/read", id: RequestId, params: GetAccountParams, } | { "method": "getConversationSummary", id: RequestId, params: GetConversationSummaryParams, } | { "method": "gitDiffToRemote", id: RequestId, params: GitDiffToRemoteParams, } | { "method": "getAuthStatus", id: RequestId, params: GetAuthStatusParams, } | { "method": "fuzzyFileSearch", id: RequestId, params: FuzzyFileSearchParams, }; diff --git a/codex-rs/app-server-protocol/schema/typescript/v2/UsageContributorKind.ts b/codex-rs/app-server-protocol/schema/typescript/v2/UsageContributorKind.ts new file mode 100644 index 0000000000..efea43b657 --- /dev/null +++ b/codex-rs/app-server-protocol/schema/typescript/v2/UsageContributorKind.ts @@ -0,0 +1,5 @@ +// GENERATED CODE! DO NOT MODIFY BY HAND! + +// This file was generated by [ts-rs](https://github.com/Aleph-Alpha/ts-rs). Do not edit this file manually. + +export type UsageContributorKind = "skill" | "subagent" | "app" | "mcpServer" | "plugin"; diff --git a/codex-rs/app-server-protocol/schema/typescript/v2/UsageEntry.ts b/codex-rs/app-server-protocol/schema/typescript/v2/UsageEntry.ts new file mode 100644 index 0000000000..d2e4486b7e --- /dev/null +++ b/codex-rs/app-server-protocol/schema/typescript/v2/UsageEntry.ts @@ -0,0 +1,6 @@ +// GENERATED CODE! DO NOT MODIFY BY HAND! + +// This file was generated by [ts-rs](https://github.com/Aleph-Alpha/ts-rs). Do not edit this file manually. +import type { UsageContributorKind } from "./UsageContributorKind"; + +export type UsageEntry = { kind: UsageContributorKind, id: string, label: string, attributedTokens: number, percentOfUsage: number, }; diff --git a/codex-rs/app-server-protocol/schema/typescript/v2/UsageHeadline.ts b/codex-rs/app-server-protocol/schema/typescript/v2/UsageHeadline.ts new file mode 100644 index 0000000000..6b110fb15f --- /dev/null +++ b/codex-rs/app-server-protocol/schema/typescript/v2/UsageHeadline.ts @@ -0,0 +1,6 @@ +// GENERATED CODE! DO NOT MODIFY BY HAND! + +// This file was generated by [ts-rs](https://github.com/Aleph-Alpha/ts-rs). Do not edit this file manually. +import type { UsageEntry } from "./UsageEntry"; + +export type UsageHeadline = { entry: UsageEntry, note: string | null, }; diff --git a/codex-rs/app-server-protocol/schema/typescript/v2/UsageRange.ts b/codex-rs/app-server-protocol/schema/typescript/v2/UsageRange.ts new file mode 100644 index 0000000000..c38e3b7184 --- /dev/null +++ b/codex-rs/app-server-protocol/schema/typescript/v2/UsageRange.ts @@ -0,0 +1,5 @@ +// GENERATED CODE! DO NOT MODIFY BY HAND! + +// This file was generated by [ts-rs](https://github.com/Aleph-Alpha/ts-rs). Do not edit this file manually. + +export type UsageRange = "day" | "week"; diff --git a/codex-rs/app-server-protocol/schema/typescript/v2/UsageReadParams.ts b/codex-rs/app-server-protocol/schema/typescript/v2/UsageReadParams.ts new file mode 100644 index 0000000000..1334673b4c --- /dev/null +++ b/codex-rs/app-server-protocol/schema/typescript/v2/UsageReadParams.ts @@ -0,0 +1,6 @@ +// GENERATED CODE! DO NOT MODIFY BY HAND! + +// This file was generated by [ts-rs](https://github.com/Aleph-Alpha/ts-rs). Do not edit this file manually. +import type { UsageRange } from "./UsageRange"; + +export type UsageReadParams = { range: UsageRange, }; diff --git a/codex-rs/app-server-protocol/schema/typescript/v2/UsageReadResponse.ts b/codex-rs/app-server-protocol/schema/typescript/v2/UsageReadResponse.ts new file mode 100644 index 0000000000..f13c8fa368 --- /dev/null +++ b/codex-rs/app-server-protocol/schema/typescript/v2/UsageReadResponse.ts @@ -0,0 +1,6 @@ +// GENERATED CODE! DO NOT MODIFY BY HAND! + +// This file was generated by [ts-rs](https://github.com/Aleph-Alpha/ts-rs). Do not edit this file manually. +import type { UsageReport } from "./UsageReport"; + +export type UsageReadResponse = { report: UsageReport, }; diff --git a/codex-rs/app-server-protocol/schema/typescript/v2/UsageReport.ts b/codex-rs/app-server-protocol/schema/typescript/v2/UsageReport.ts new file mode 100644 index 0000000000..24d776fc2e --- /dev/null +++ b/codex-rs/app-server-protocol/schema/typescript/v2/UsageReport.ts @@ -0,0 +1,8 @@ +// GENERATED CODE! DO NOT MODIFY BY HAND! + +// This file was generated by [ts-rs](https://github.com/Aleph-Alpha/ts-rs). Do not edit this file manually. +import type { UsageEntry } from "./UsageEntry"; +import type { UsageHeadline } from "./UsageHeadline"; +import type { UsageRange } from "./UsageRange"; + +export type UsageReport = { range: UsageRange, generatedAt: number, trackedFrom: number | null, totalTokens: number, headline: UsageHeadline | null, skills: Array, subagents: Array, apps: Array, mcpServers: Array, plugins: Array, }; diff --git a/codex-rs/app-server-protocol/schema/typescript/v2/index.ts b/codex-rs/app-server-protocol/schema/typescript/v2/index.ts index e9b8993c26..3b01a3f1c9 100644 --- a/codex-rs/app-server-protocol/schema/typescript/v2/index.ts +++ b/codex-rs/app-server-protocol/schema/typescript/v2/index.ts @@ -448,6 +448,13 @@ export type { TurnStartedNotification } from "./TurnStartedNotification"; export type { TurnStatus } from "./TurnStatus"; export type { TurnSteerParams } from "./TurnSteerParams"; export type { TurnSteerResponse } from "./TurnSteerResponse"; +export type { UsageContributorKind } from "./UsageContributorKind"; +export type { UsageEntry } from "./UsageEntry"; +export type { UsageHeadline } from "./UsageHeadline"; +export type { UsageRange } from "./UsageRange"; +export type { UsageReadParams } from "./UsageReadParams"; +export type { UsageReadResponse } from "./UsageReadResponse"; +export type { UsageReport } from "./UsageReport"; export type { UserInput } from "./UserInput"; export type { WarningNotification } from "./WarningNotification"; export type { WebSearchAction } from "./WebSearchAction"; diff --git a/codex-rs/app-server-protocol/src/protocol/common.rs b/codex-rs/app-server-protocol/src/protocol/common.rs index 76334e8a6a..ffde24c977 100644 --- a/codex-rs/app-server-protocol/src/protocol/common.rs +++ b/codex-rs/app-server-protocol/src/protocol/common.rs @@ -629,6 +629,11 @@ client_request_definitions! { serialization: None, response: v2::PluginListResponse, }, + UsageRead => "usage/read" { + params: v2::UsageReadParams, + serialization: None, + response: v2::UsageReadResponse, + }, PluginInstalled => "plugin/installed" { params: v2::PluginInstalledParams, serialization: None, diff --git a/codex-rs/app-server-protocol/src/protocol/thread_history.rs b/codex-rs/app-server-protocol/src/protocol/thread_history.rs index 0fbb93f397..467abf3fd3 100644 --- a/codex-rs/app-server-protocol/src/protocol/thread_history.rs +++ b/codex-rs/app-server-protocol/src/protocol/thread_history.rs @@ -232,7 +232,9 @@ impl ThreadHistoryBuilder { RolloutItem::EventMsg(event) => self.handle_event(event), RolloutItem::Compacted(payload) => self.handle_compacted(payload), RolloutItem::ResponseItem(item) => self.handle_response_item(item), - RolloutItem::TurnContext(_) | RolloutItem::SessionMeta(_) => {} + RolloutItem::TurnContext(_) + | RolloutItem::SessionMeta(_) + | RolloutItem::UsageAttribution(_) => {} } } diff --git a/codex-rs/app-server-protocol/src/protocol/v2/mod.rs b/codex-rs/app-server-protocol/src/protocol/v2/mod.rs index b5fa9fdc65..1ac1f898cf 100644 --- a/codex-rs/app-server-protocol/src/protocol/v2/mod.rs +++ b/codex-rs/app-server-protocol/src/protocol/v2/mod.rs @@ -24,6 +24,7 @@ mod review; mod thread; mod thread_data; mod turn; +mod usage; mod windows_sandbox; pub use account::*; @@ -51,6 +52,7 @@ pub use shared::*; pub use thread::*; pub use thread_data::*; pub use turn::*; +pub use usage::*; pub use windows_sandbox::*; #[cfg(test)] diff --git a/codex-rs/app-server-protocol/src/protocol/v2/usage.rs b/codex-rs/app-server-protocol/src/protocol/v2/usage.rs new file mode 100644 index 0000000000..b6592609f6 --- /dev/null +++ b/codex-rs/app-server-protocol/src/protocol/v2/usage.rs @@ -0,0 +1,89 @@ +use codex_protocol::protocol::UsageContributorKind as CoreUsageContributorKind; +use schemars::JsonSchema; +use serde::Deserialize; +use serde::Serialize; +use ts_rs::TS; + +#[derive(Serialize, Deserialize, Debug, Clone, Copy, PartialEq, Eq, JsonSchema, TS)] +#[serde(rename_all = "camelCase")] +#[ts(rename_all = "camelCase", export_to = "v2/")] +pub enum UsageRange { + Day, + Week, +} + +#[derive(Serialize, Deserialize, Debug, Clone, PartialEq, Eq, JsonSchema, TS)] +#[serde(rename_all = "camelCase")] +#[ts(export_to = "v2/")] +pub struct UsageReadParams { + pub range: UsageRange, +} + +#[derive(Serialize, Deserialize, Debug, Clone, PartialEq, Eq, JsonSchema, TS)] +#[serde(rename_all = "camelCase")] +#[ts(export_to = "v2/")] +pub struct UsageEntry { + pub kind: UsageContributorKind, + pub id: String, + pub label: String, + #[ts(type = "number")] + pub attributed_tokens: i64, + pub percent_of_usage: u8, +} + +#[derive(Serialize, Deserialize, Debug, Clone, PartialEq, Eq, JsonSchema, TS)] +#[serde(rename_all = "camelCase")] +#[ts(export_to = "v2/")] +pub struct UsageHeadline { + pub entry: UsageEntry, + pub note: Option, +} + +#[derive(Serialize, Deserialize, Debug, Clone, PartialEq, Eq, JsonSchema, TS)] +#[serde(rename_all = "camelCase")] +#[ts(export_to = "v2/")] +pub struct UsageReport { + pub range: UsageRange, + #[ts(type = "number")] + pub generated_at: i64, + #[ts(type = "number | null")] + pub tracked_from: Option, + #[ts(type = "number")] + pub total_tokens: i64, + pub headline: Option, + pub skills: Vec, + pub subagents: Vec, + pub apps: Vec, + pub mcp_servers: Vec, + pub plugins: Vec, +} + +#[derive(Serialize, Deserialize, Debug, Clone, PartialEq, Eq, JsonSchema, TS)] +#[serde(rename_all = "camelCase")] +#[ts(export_to = "v2/")] +pub struct UsageReadResponse { + pub report: UsageReport, +} + +#[derive(Serialize, Deserialize, Debug, Clone, Copy, PartialEq, Eq, JsonSchema, TS)] +#[serde(rename_all = "camelCase")] +#[ts(rename_all = "camelCase", export_to = "v2/")] +pub enum UsageContributorKind { + Skill, + Subagent, + App, + McpServer, + Plugin, +} + +impl From for UsageContributorKind { + fn from(value: CoreUsageContributorKind) -> Self { + match value { + CoreUsageContributorKind::Skill => Self::Skill, + CoreUsageContributorKind::Subagent => Self::Subagent, + CoreUsageContributorKind::App => Self::App, + CoreUsageContributorKind::McpServer => Self::McpServer, + CoreUsageContributorKind::Plugin => Self::Plugin, + } + } +} diff --git a/codex-rs/app-server/README.md b/codex-rs/app-server/README.md index 2ceffc86fe..ef14008eed 100644 --- a/codex-rs/app-server/README.md +++ b/codex-rs/app-server/README.md @@ -206,6 +206,7 @@ Example with notification opt-out: - `plugin/skill/read` — read remote plugin skill markdown on demand by `remoteMarketplaceName`, `remotePluginId`, and `skillName`. This lets clients preview uninstalled remote plugin skills without downloading the plugin bundle. - `skills/changed` — notification emitted when watched local skill files change. - `app/list` — list available apps. +- `usage/read` — read forward-only local token usage for a rolling `day` or `week` range, grouped by skills, subagents, apps, MCP servers, and plugins when Codex has tracked those contributors in sqlite. - `remoteControl/enable` — experimental; enable remote control for the current app-server process and return the current remote-control status snapshot. The caller is responsible for persisting the desired setting outside app-server. - `remoteControl/disable` — experimental; disable remote control for the current app-server process and return the current remote-control status snapshot. This does not revoke already enrolled controller devices. - `remoteControl/status/read` — experimental; read the current remote-control status snapshot. `status` is one of `disabled`, `connecting`, `connected`, or `errored`; `serverName` is the local machine name used by this app-server process; `environmentId` is a string when the app-server has a current enrollment and `null` when that enrollment is cleared, invalidated, or remote control is disabled. diff --git a/codex-rs/app-server/src/message_processor.rs b/codex-rs/app-server/src/message_processor.rs index 9c10a33408..0a75233486 100644 --- a/codex-rs/app-server/src/message_processor.rs +++ b/codex-rs/app-server/src/message_processor.rs @@ -35,6 +35,7 @@ use crate::request_processors::SearchRequestProcessor; use crate::request_processors::ThreadGoalRequestProcessor; use crate::request_processors::ThreadRequestProcessor; use crate::request_processors::TurnRequestProcessor; +use crate::request_processors::UsageRequestProcessor; use crate::request_processors::WindowsSandboxRequestProcessor; use crate::request_serialization::QueuedInitializedRequest; use crate::request_serialization::RequestSerializationQueueKey; @@ -181,6 +182,7 @@ pub(crate) struct MessageProcessor { thread_goal_processor: ThreadGoalRequestProcessor, thread_processor: ThreadRequestProcessor, turn_processor: TurnRequestProcessor, + usage_processor: UsageRequestProcessor, windows_sandbox_processor: WindowsSandboxRequestProcessor, request_serialization_queues: RequestSerializationQueues, } @@ -423,9 +425,10 @@ impl MessageProcessor { thread_watch_manager.clone(), Arc::clone(&thread_list_state_permit), thread_goal_processor.clone(), - state_db, + state_db.clone(), Arc::clone(&skills_watcher), ); + let usage_processor = UsageRequestProcessor::new(state_db); let turn_processor = TurnRequestProcessor::new( auth_manager.clone(), Arc::clone(&thread_manager), @@ -502,6 +505,7 @@ impl MessageProcessor { thread_goal_processor, thread_processor, turn_processor, + usage_processor, windows_sandbox_processor, request_serialization_queues: RequestSerializationQueues::default(), } @@ -1115,6 +1119,11 @@ impl MessageProcessor { ClientRequest::PluginList { params, .. } => { self.plugin_processor.plugin_list(params).await } + ClientRequest::UsageRead { params, .. } => self + .usage_processor + .usage_read(params) + .await + .map(|response| Some(response.into())), ClientRequest::PluginInstalled { params, .. } => { self.plugin_processor.plugin_installed(params).await } diff --git a/codex-rs/app-server/src/request_processors.rs b/codex-rs/app-server/src/request_processors.rs index 838dd8adf4..4c588a2d8f 100644 --- a/codex-rs/app-server/src/request_processors.rs +++ b/codex-rs/app-server/src/request_processors.rs @@ -467,6 +467,7 @@ mod search; mod thread_processor; mod token_usage_replay; mod turn_processor; +mod usage_processor; mod windows_sandbox_processor; pub(crate) use account_processor::AccountRequestProcessor; @@ -489,6 +490,7 @@ pub(crate) use search::SearchRequestProcessor; pub(crate) use thread_goal_processor::ThreadGoalRequestProcessor; pub(crate) use thread_processor::ThreadRequestProcessor; pub(crate) use turn_processor::TurnRequestProcessor; +pub(crate) use usage_processor::UsageRequestProcessor; pub(crate) use windows_sandbox_processor::WindowsSandboxRequestProcessor; use crate::error_code::internal_error; diff --git a/codex-rs/app-server/src/request_processors/usage_processor.rs b/codex-rs/app-server/src/request_processors/usage_processor.rs new file mode 100644 index 0000000000..aefa1533dc --- /dev/null +++ b/codex-rs/app-server/src/request_processors/usage_processor.rs @@ -0,0 +1,76 @@ +use super::*; +use crate::error_code::internal_error; +use chrono::Utc; +use codex_app_server_protocol::UsageEntry; +use codex_app_server_protocol::UsageHeadline; +use codex_app_server_protocol::UsageRange; +use codex_app_server_protocol::UsageReadParams; +use codex_app_server_protocol::UsageReadResponse; +use codex_app_server_protocol::UsageReport; +use codex_rollout::StateDbHandle; + +#[derive(Clone)] +pub(crate) struct UsageRequestProcessor { + state_db: Option, +} + +impl UsageRequestProcessor { + pub(crate) fn new(state_db: Option) -> Self { + Self { state_db } + } + + pub(crate) async fn usage_read( + &self, + params: UsageReadParams, + ) -> Result { + let state_db = self + .state_db + .as_ref() + .ok_or_else(|| internal_error("sqlite state db unavailable for usage"))?; + let report = state_db + .read_usage_report(state_usage_range(params.range), Utc::now().timestamp()) + .await + .map_err(|err| internal_error(format!("failed to read usage report: {err}")))?; + Ok(UsageReadResponse { + report: UsageReport { + range: api_usage_range(report.range), + generated_at: report.generated_at, + tracked_from: report.tracked_from, + total_tokens: report.total_tokens, + headline: report.headline.map(|headline| UsageHeadline { + entry: usage_entry(headline.entry), + note: headline.note, + }), + skills: report.skills.into_iter().map(usage_entry).collect(), + subagents: report.subagents.into_iter().map(usage_entry).collect(), + apps: report.apps.into_iter().map(usage_entry).collect(), + mcp_servers: report.mcp_servers.into_iter().map(usage_entry).collect(), + plugins: report.plugins.into_iter().map(usage_entry).collect(), + }, + }) + } +} + +fn usage_entry(entry: codex_state::UsageEntry) -> UsageEntry { + UsageEntry { + kind: entry.kind.into(), + id: entry.id, + label: entry.label, + attributed_tokens: entry.attributed_tokens, + percent_of_usage: entry.percent_of_usage, + } +} + +fn state_usage_range(value: UsageRange) -> codex_state::UsageRange { + match value { + UsageRange::Day => codex_state::UsageRange::Day, + UsageRange::Week => codex_state::UsageRange::Week, + } +} + +fn api_usage_range(value: codex_state::UsageRange) -> UsageRange { + match value { + codex_state::UsageRange::Day => UsageRange::Day, + codex_state::UsageRange::Week => UsageRange::Week, + } +} diff --git a/codex-rs/core/src/agent/control.rs b/codex-rs/core/src/agent/control.rs index 42981a5d0a..8fb374237e 100644 --- a/codex-rs/core/src/agent/control.rs +++ b/codex-rs/core/src/agent/control.rs @@ -125,6 +125,7 @@ fn keep_forked_rollout_item(item: &RolloutItem, preserve_reference_context_item: // so they must rebuild context on their first child turn. RolloutItem::TurnContext(_) => preserve_reference_context_item, RolloutItem::Compacted(_) | RolloutItem::EventMsg(_) | RolloutItem::SessionMeta(_) => true, + RolloutItem::UsageAttribution(_) => false, } } diff --git a/codex-rs/core/src/client_common.rs b/codex-rs/core/src/client_common.rs index d9edadbed4..d83d92809f 100644 --- a/codex-rs/core/src/client_common.rs +++ b/codex-rs/core/src/client_common.rs @@ -1,3 +1,4 @@ +use crate::usage::UsagePromptAttribution; pub use codex_api::ResponseEvent; use codex_config::types::Personality; use codex_protocol::error::Result; @@ -33,6 +34,8 @@ pub struct Prompt { /// Whether parallel tool calls are permitted for this prompt. pub(crate) parallel_tool_calls: bool, + pub(crate) usage_attribution: UsagePromptAttribution, + pub base_instructions: BaseInstructions, /// Optionally specify the personality of the model. @@ -51,6 +54,7 @@ impl Default for Prompt { input: Vec::new(), tools: Vec::new(), parallel_tool_calls: false, + usage_attribution: UsagePromptAttribution::default(), base_instructions: BaseInstructions::default(), personality: None, output_schema: None, diff --git a/codex-rs/core/src/compact_remote.rs b/codex-rs/core/src/compact_remote.rs index c7ba1a314f..e7e427613e 100644 --- a/codex-rs/core/src/compact_remote.rs +++ b/codex-rs/core/src/compact_remote.rs @@ -180,6 +180,7 @@ async fn run_remote_compact_task_inner_impl( input: prompt_input, tools: tool_router.model_visible_specs(), parallel_tool_calls: turn_context.model_info.supports_parallel_tool_calls, + usage_attribution: Default::default(), base_instructions, personality: turn_context.personality, output_schema: None, diff --git a/codex-rs/core/src/compact_remote_v2.rs b/codex-rs/core/src/compact_remote_v2.rs index 2e261b2e94..c1d0cfe688 100644 --- a/codex-rs/core/src/compact_remote_v2.rs +++ b/codex-rs/core/src/compact_remote_v2.rs @@ -188,6 +188,7 @@ async fn run_remote_compact_task_inner_impl( input, tools: tool_router.model_visible_specs(), parallel_tool_calls: turn_context.model_info.supports_parallel_tool_calls, + usage_attribution: Default::default(), base_instructions, personality: turn_context.personality, output_schema: None, diff --git a/codex-rs/core/src/lib.rs b/codex-rs/core/src/lib.rs index d5c1237638..2081c3257d 100644 --- a/codex-rs/core/src/lib.rs +++ b/codex-rs/core/src/lib.rs @@ -98,6 +98,7 @@ pub(crate) use skills::skills_load_input_from_config; mod stream_events_utils; pub mod test_support; mod unified_exec; +mod usage; pub mod windows_sandbox; pub use client::X_RESPONSESAPI_INCLUDE_TIMING_METRICS_HEADER; pub use codex_protocol::config_types::ModelProviderAuthInfo; diff --git a/codex-rs/core/src/session/mod.rs b/codex-rs/core/src/session/mod.rs index c12c29e7e1..2ff89afd7d 100644 --- a/codex-rs/core/src/session/mod.rs +++ b/codex-rs/core/src/session/mod.rs @@ -3001,6 +3001,38 @@ impl Session { } } + pub(crate) async fn record_usage_attribution( + &self, + turn_context: &TurnContext, + prompt: &crate::client_common::Prompt, + response_id: &str, + token_usage: Option<&TokenUsage>, + ) { + let Some(token_usage) = token_usage else { + return; + }; + let occurred_at = chrono::Utc::now().timestamp(); + let attribution = prompt.usage_attribution.complete( + format!("{}:{response_id}", self.conversation_id), + turn_context.sub_id.clone(), + response_id.to_string(), + occurred_at, + token_usage.clone(), + ); + self.persist_rollout_items(&[RolloutItem::UsageAttribution(attribution.clone())]) + .await; + if let Some(state_db) = self.state_db() + && let Err(err) = state_db + .record_usage_sample(&codex_state::UsageSample { + thread_id: self.conversation_id, + attribution, + }) + .await + { + warn!("failed to persist usage sample: {err}"); + } + } + pub(crate) async fn recompute_token_usage(&self, turn_context: &TurnContext) { let history = self.clone_history().await; let base_instructions = self.get_base_instructions().await; diff --git a/codex-rs/core/src/session/rollout_reconstruction.rs b/codex-rs/core/src/session/rollout_reconstruction.rs index a4c042af0c..e66dac03ce 100644 --- a/codex-rs/core/src/session/rollout_reconstruction.rs +++ b/codex-rs/core/src/session/rollout_reconstruction.rs @@ -207,7 +207,9 @@ impl Session { active_segment.get_or_insert_with(ActiveReplaySegment::default); active_segment.counts_as_user_turn |= is_user_turn_boundary(response_item); } - RolloutItem::EventMsg(_) | RolloutItem::SessionMeta(_) => {} + RolloutItem::EventMsg(_) + | RolloutItem::SessionMeta(_) + | RolloutItem::UsageAttribution(_) => {} } if base_replacement_history.is_some() @@ -276,7 +278,8 @@ impl Session { } RolloutItem::EventMsg(_) | RolloutItem::TurnContext(_) - | RolloutItem::SessionMeta(_) => {} + | RolloutItem::SessionMeta(_) + | RolloutItem::UsageAttribution(_) => {} } } diff --git a/codex-rs/core/src/session/turn.rs b/codex-rs/core/src/session/turn.rs index 26ba845460..195282296b 100644 --- a/codex-rs/core/src/session/turn.rs +++ b/codex-rs/core/src/session/turn.rs @@ -889,10 +889,16 @@ pub(crate) fn build_prompt( turn_context: &TurnContext, base_instructions: BaseInstructions, ) -> Prompt { + let usage_attribution = crate::usage::UsagePromptAttribution::from_prompt( + &input, + router, + base_instructions.text.as_str(), + ); Prompt { input, tools: router.model_visible_specs(), parallel_tool_calls: turn_context.model_info.supports_parallel_tool_calls, + usage_attribution, base_instructions, personality: turn_context.personality, output_schema: turn_context.final_output_json_schema.clone(), @@ -2021,6 +2027,13 @@ async fn try_run_sampling_request( &mut assistant_message_stream_parsers, ) .await; + sess.record_usage_attribution( + &turn_context, + prompt, + response_id.as_str(), + token_usage.as_ref(), + ) + .await; sess.record_token_usage_info(&turn_context, token_usage.as_ref()) .await; should_emit_token_count = true; diff --git a/codex-rs/core/src/tools/handlers/mcp.rs b/codex-rs/core/src/tools/handlers/mcp.rs index 6f4f79505b..ee466dffed 100644 --- a/codex-rs/core/src/tools/handlers/mcp.rs +++ b/codex-rs/core/src/tools/handlers/mcp.rs @@ -18,6 +18,8 @@ use crate::tools::registry::ToolExposure; use crate::tools::registry::ToolTelemetryTags; use crate::tools::tool_search_entry::ToolSearchInfo; use codex_mcp::ToolInfo; +use codex_protocol::protocol::UsageContributor; +use codex_protocol::protocol::UsageContributorKind; use codex_tools::ResponsesApiNamespace; use codex_tools::ResponsesApiNamespaceTool; use codex_tools::ToolName; @@ -131,6 +133,38 @@ impl ToolExecutor for McpHandler { } impl CoreToolRuntime for McpHandler { + fn usage_contributors(&self) -> Vec { + let mut contributors = Vec::new(); + if let Some(connector_id) = self.tool_info.connector_id.as_ref() { + contributors.push(UsageContributor { + kind: UsageContributorKind::App, + id: connector_id.clone(), + label: self + .tool_info + .connector_name + .clone() + .unwrap_or_else(|| connector_id.clone()), + }); + } else { + contributors.push(UsageContributor { + kind: UsageContributorKind::McpServer, + id: self.tool_info.server_name.clone(), + label: self.tool_info.server_name.clone(), + }); + } + contributors.extend( + self.tool_info + .plugin_display_names + .iter() + .map(|plugin_name| UsageContributor { + kind: UsageContributorKind::Plugin, + id: plugin_name.clone(), + label: plugin_name.clone(), + }), + ); + contributors + } + fn search_info(&self) -> Option { let source_name = self .tool_info diff --git a/codex-rs/core/src/tools/parallel.rs b/codex-rs/core/src/tools/parallel.rs index 15954869e2..41d98c5994 100644 --- a/codex-rs/core/src/tools/parallel.rs +++ b/codex-rs/core/src/tools/parallel.rs @@ -320,6 +320,8 @@ mod tests { let router = Arc::new(ToolRouter::from_parts( ToolRegistry::from_tools([handler]), Vec::new(), + Vec::new(), + std::collections::HashMap::new(), )); let tracker = Arc::new(tokio::sync::Mutex::new(TurnDiffTracker::new())); let runtime = ToolCallRuntime::new(router, session, turn_context, tracker); diff --git a/codex-rs/core/src/tools/registry.rs b/codex-rs/core/src/tools/registry.rs index 96837497f8..c4415ff3dd 100644 --- a/codex-rs/core/src/tools/registry.rs +++ b/codex-rs/core/src/tools/registry.rs @@ -28,6 +28,7 @@ use crate::util::error_or_panic; use codex_extension_api::ToolCallOutcome; use codex_protocol::models::ResponseInputItem; use codex_protocol::protocol::EventMsg; +use codex_protocol::protocol::UsageContributor; use codex_tools::ToolName; use codex_tools::ToolSpec; use futures::future::BoxFuture; @@ -70,6 +71,10 @@ pub(crate) trait CoreToolRuntime: ToolExecutor { None } + fn usage_contributors(&self) -> Vec { + Vec::new() + } + fn pre_tool_use_payload(&self, _invocation: &ToolInvocation) -> Option { None } diff --git a/codex-rs/core/src/tools/router.rs b/codex-rs/core/src/tools/router.rs index cfd4776890..939b3ea371 100644 --- a/codex-rs/core/src/tools/router.rs +++ b/codex-rs/core/src/tools/router.rs @@ -17,6 +17,7 @@ use codex_tools::ToolCall as ExtensionToolCall; use codex_tools::ToolExecutor; use codex_tools::ToolName; use codex_tools::ToolSpec; +use std::collections::HashMap; use std::sync::Arc; use std::sync::atomic::AtomicBool; use tokio_util::sync::CancellationToken; @@ -34,6 +35,9 @@ pub struct ToolCall { pub struct ToolRouter { registry: ToolRegistry, model_visible_specs: Vec, + usage_contributors: Vec, + usage_contributors_by_tool_name: + HashMap>, } pub(crate) struct ToolRouterParams<'a> { @@ -49,10 +53,20 @@ impl ToolRouter { build_tool_router(turn_context, params) } - pub(crate) fn from_parts(registry: ToolRegistry, model_visible_specs: Vec) -> Self { + pub(crate) fn from_parts( + registry: ToolRegistry, + model_visible_specs: Vec, + usage_contributors: Vec, + usage_contributors_by_tool_name: HashMap< + ToolName, + Vec, + >, + ) -> Self { Self { registry, model_visible_specs, + usage_contributors, + usage_contributors_by_tool_name, } } @@ -60,6 +74,20 @@ impl ToolRouter { self.model_visible_specs.clone() } + pub(crate) fn usage_contributors(&self) -> Vec { + self.usage_contributors.clone() + } + + pub(crate) fn usage_contributors_for_tool_name( + &self, + tool_name: &ToolName, + ) -> Vec { + self.usage_contributors_by_tool_name + .get(tool_name) + .cloned() + .unwrap_or_default() + } + #[cfg(test)] pub(crate) fn registered_tool_names_for_test(&self) -> Vec { self.registry.tool_names_for_test() diff --git a/codex-rs/core/src/tools/spec_plan.rs b/codex-rs/core/src/tools/spec_plan.rs index a47ed9da34..32b0d592f8 100644 --- a/codex-rs/core/src/tools/spec_plan.rs +++ b/codex-rs/core/src/tools/spec_plan.rs @@ -76,6 +76,7 @@ use codex_tools::request_user_input_available_modes; use codex_tools::shell_command_backend_for_features; use codex_tools::shell_type_for_model_and_features; use std::collections::BTreeMap; +use std::collections::HashMap; use std::collections::HashSet; use std::sync::Arc; use tracing::warn; @@ -123,14 +124,25 @@ pub(crate) fn build_tool_router( turn_context: &TurnContext, params: ToolRouterParams<'_>, ) -> ToolRouter { - let (model_visible_specs, registry) = build_tool_specs_and_registry(turn_context, params); - ToolRouter::from_parts(registry, model_visible_specs) + let (model_visible_specs, registry, usage_contributors, usage_contributors_by_tool_name) = + build_tool_specs_and_registry(turn_context, params); + ToolRouter::from_parts( + registry, + model_visible_specs, + usage_contributors, + usage_contributors_by_tool_name, + ) } fn build_tool_specs_and_registry( turn_context: &TurnContext, params: ToolRouterParams<'_>, -) -> (Vec, ToolRegistry) { +) -> ( + Vec, + ToolRegistry, + Vec, + HashMap>, +) { let ToolRouterParams { mcp_tools, deferred_mcp_tools, @@ -160,12 +172,19 @@ fn build_tool_specs_and_registry( fn build_model_visible_specs_and_registry( turn_context: &TurnContext, planned_tools: PlannedTools, -) -> (Vec, ToolRegistry) { +) -> ( + Vec, + ToolRegistry, + Vec, + HashMap>, +) { let PlannedTools { runtimes, hosted_specs, } = planned_tools; let mut specs = Vec::new(); + let mut usage_contributors = Vec::new(); + let mut usage_contributors_by_tool_name = HashMap::new(); let mut seen_tool_names = HashSet::new(); for runtime in &runtimes { let tool_name = runtime.tool_name(); @@ -177,6 +196,16 @@ fn build_model_visible_specs_and_registry( && !is_hidden_by_code_mode_only(turn_context, &tool_name, exposure) && let Some(spec) = runtime.spec() { + let estimated_tokens = crate::usage::estimate_serialized_tokens(&spec); + let runtime_usage_contributors = runtime.usage_contributors(); + usage_contributors_by_tool_name + .insert(tool_name.clone(), runtime_usage_contributors.clone()); + usage_contributors.extend(runtime_usage_contributors.into_iter().map(|contributor| { + crate::usage::UsagePromptContributor { + contributor, + source_estimated_tokens: estimated_tokens, + } + })); specs.push(spec_for_model_request(turn_context, exposure, spec)); } } @@ -198,7 +227,12 @@ fn build_model_visible_specs_and_registry( }) .collect(); - (model_visible_specs, registry) + ( + model_visible_specs, + registry, + usage_contributors, + usage_contributors_by_tool_name, + ) } fn spec_for_model_request( diff --git a/codex-rs/core/src/usage.rs b/codex-rs/core/src/usage.rs new file mode 100644 index 0000000000..df6ec1d87a --- /dev/null +++ b/codex-rs/core/src/usage.rs @@ -0,0 +1,363 @@ +use crate::tools::router::ToolRouter; +use codex_protocol::models::ContentItem; +use codex_protocol::models::ResponseItem; +use codex_protocol::protocol::TokenUsage; +use codex_protocol::protocol::UsageAttributionContributor; +use codex_protocol::protocol::UsageAttributionItem; +use codex_protocol::protocol::UsageContributor; +use codex_protocol::protocol::UsageContributorKind; +use codex_tools::ToolName; +use codex_utils_output_truncation::approx_token_count; +use std::collections::BTreeMap; +use std::collections::HashMap; + +#[derive(Debug, Clone, Default, PartialEq, Eq)] +pub(crate) struct UsagePromptAttribution { + pub(crate) prompt_estimated_tokens: i64, + pub(crate) contributors: Vec, +} + +#[derive(Debug, Clone, PartialEq, Eq)] +pub(crate) struct UsagePromptContributor { + pub(crate) contributor: UsageContributor, + pub(crate) source_estimated_tokens: i64, +} + +impl UsagePromptAttribution { + pub(crate) fn from_prompt( + input: &[ResponseItem], + router: &ToolRouter, + base_instructions: &str, + ) -> Self { + let mut contributors = skill_contributors(input); + contributors.extend(router.usage_contributors()); + contributors.extend(tool_result_contributors(input, router)); + let input_tokens = input + .iter() + .map(estimate_response_item_tokens) + .fold(0i64, i64::saturating_add); + let tool_tokens = router + .model_visible_specs() + .iter() + .map(estimate_serialized_tokens) + .fold(0i64, i64::saturating_add); + let base_tokens = i64::try_from(approx_token_count(base_instructions)).unwrap_or(i64::MAX); + Self { + prompt_estimated_tokens: base_tokens + .saturating_add(input_tokens) + .saturating_add(tool_tokens), + contributors: aggregate_contributors(contributors), + } + } + + pub(crate) fn complete( + &self, + sample_id: String, + turn_id: String, + response_id: String, + occurred_at: i64, + token_usage: TokenUsage, + ) -> UsageAttributionItem { + let non_cached_input = token_usage.non_cached_input(); + let contributors = self + .contributors + .iter() + .map(|contributor| UsageAttributionContributor { + contributor: contributor.contributor.clone(), + source_estimated_tokens: contributor.source_estimated_tokens, + attributed_tokens: attributable_tokens( + non_cached_input, + contributor.source_estimated_tokens, + self.prompt_estimated_tokens, + ), + }) + .filter(|contributor| contributor.attributed_tokens > 0) + .collect(); + UsageAttributionItem { + sample_id, + turn_id, + response_id, + occurred_at, + token_usage, + prompt_estimated_tokens: self.prompt_estimated_tokens, + contributors, + } + } +} + +pub(crate) fn estimate_serialized_tokens(value: &T) -> i64 { + serde_json::to_string(value) + .map(|serialized| i64::try_from(approx_token_count(&serialized)).unwrap_or(i64::MAX)) + .unwrap_or(/*default*/ 0) +} + +fn estimate_response_item_tokens(item: &ResponseItem) -> i64 { + estimate_serialized_tokens(item) +} + +fn skill_contributors(input: &[ResponseItem]) -> Vec { + input.iter().filter_map(skill_contributor).collect() +} + +fn tool_result_contributors( + input: &[ResponseItem], + router: &ToolRouter, +) -> Vec { + let contributors_by_call_id = input + .iter() + .filter_map(|item| tool_call_contributors(item, router)) + .collect::>(); + input + .iter() + .filter_map(|item| tool_result_contributor(item, &contributors_by_call_id)) + .flatten() + .collect() +} + +fn tool_call_contributors( + item: &ResponseItem, + router: &ToolRouter, +) -> Option<(String, Vec)> { + let (call_id, tool_name) = match item { + ResponseItem::FunctionCall { + call_id, + name, + namespace, + .. + } => (call_id, ToolName::new(namespace.clone(), name)), + ResponseItem::CustomToolCall { call_id, name, .. } => (call_id, ToolName::plain(name)), + _ => return None, + }; + let contributors = router.usage_contributors_for_tool_name(&tool_name); + (!contributors.is_empty()).then(|| (call_id.clone(), contributors)) +} + +fn tool_result_contributor( + item: &ResponseItem, + contributors_by_call_id: &HashMap>, +) -> Option> { + let call_id = match item { + ResponseItem::FunctionCallOutput { call_id, .. } + | ResponseItem::CustomToolCallOutput { call_id, .. } => call_id, + _ => return None, + }; + let source_estimated_tokens = estimate_response_item_tokens(item); + Some( + contributors_by_call_id + .get(call_id)? + .iter() + .cloned() + .map(|contributor| UsagePromptContributor { + contributor, + source_estimated_tokens, + }) + .collect(), + ) +} + +fn skill_contributor(item: &ResponseItem) -> Option { + let ResponseItem::Message { content, .. } = item else { + return None; + }; + let text = content.iter().find_map(|content| match content { + ContentItem::InputText { text } if text.contains("") => Some(text.as_str()), + _ => None, + })?; + let name = tag_contents(text, "name")?; + let path = tag_contents(text, "path")?; + Some(UsagePromptContributor { + contributor: UsageContributor { + kind: UsageContributorKind::Skill, + id: path.to_string(), + label: name.to_string(), + }, + source_estimated_tokens: i64::try_from(approx_token_count(text)).unwrap_or(i64::MAX), + }) +} + +fn tag_contents<'a>(text: &'a str, tag: &str) -> Option<&'a str> { + let open = format!("<{tag}>"); + let close = format!(""); + let start = text.find(open.as_str())? + open.len(); + let end = text[start..].find(close.as_str())? + start; + Some(text[start..end].trim()) +} + +fn aggregate_contributors( + contributors: Vec, +) -> Vec { + let mut aggregated = BTreeMap::new(); + for contributor in contributors { + let key = ( + contributor.contributor.kind as u8, + contributor.contributor.id.clone(), + contributor.contributor.label.clone(), + ); + aggregated + .entry(key) + .and_modify(|existing: &mut UsagePromptContributor| { + existing.source_estimated_tokens = existing + .source_estimated_tokens + .saturating_add(contributor.source_estimated_tokens); + }) + .or_insert(contributor); + } + aggregated.into_values().collect() +} + +fn attributable_tokens(non_cached_input: i64, source_tokens: i64, prompt_tokens: i64) -> i64 { + if non_cached_input <= 0 || source_tokens <= 0 || prompt_tokens <= 0 { + return 0; + } + non_cached_input + .saturating_mul(source_tokens) + .saturating_add(prompt_tokens / 2) + / prompt_tokens +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::tools::registry::ToolRegistry; + use codex_protocol::models::FunctionCallOutputPayload; + use pretty_assertions::assert_eq; + + #[test] + fn complete_attributes_only_non_cached_input_tokens() { + let attribution = UsagePromptAttribution { + prompt_estimated_tokens: 100, + contributors: vec![ + usage_prompt_contributor( + UsageContributorKind::Skill, + "/skills/tmux", + "tmux", + /*source_estimated_tokens*/ 25, + ), + usage_prompt_contributor( + UsageContributorKind::App, + "slack", + "Slack", + /*source_estimated_tokens*/ 10, + ), + ], + }; + + let usage = attribution.complete( + "sample".to_string(), + "turn".to_string(), + "response".to_string(), + /*occurred_at*/ 1_700_000_000, + TokenUsage { + input_tokens: 100, + cached_input_tokens: 40, + output_tokens: 20, + reasoning_output_tokens: 0, + total_tokens: 120, + }, + ); + + assert_eq!( + usage.contributors, + vec![ + UsageAttributionContributor { + contributor: usage_contributor( + UsageContributorKind::Skill, + "/skills/tmux", + "tmux", + ), + source_estimated_tokens: 25, + attributed_tokens: 15, + }, + UsageAttributionContributor { + contributor: usage_contributor(UsageContributorKind::App, "slack", "Slack"), + source_estimated_tokens: 10, + attributed_tokens: 6, + }, + ] + ); + } + + #[test] + fn skill_contributors_use_skill_path_as_stable_id() { + let item = ResponseItem::Message { + id: None, + role: "developer".to_string(), + content: vec![ContentItem::InputText { + text: "tmux/skills/tmux/SKILL.md" + .to_string(), + }], + phase: None, + }; + + assert_eq!( + skill_contributors(&[item]), + vec![UsagePromptContributor { + contributor: usage_contributor( + UsageContributorKind::Skill, + "/skills/tmux/SKILL.md", + "tmux", + ), + source_estimated_tokens: i64::try_from(approx_token_count( + "tmux/skills/tmux/SKILL.md", + )) + .expect("skill prompt token estimate should fit in i64"), + }] + ); + } + + #[test] + fn tool_results_reuse_tool_usage_provenance() { + let contributor = usage_contributor(UsageContributorKind::App, "slack", "Slack"); + let tool_name = ToolName::plain("mcp__slack__search"); + let router = ToolRouter::from_parts( + ToolRegistry::from_tools(Vec::< + std::sync::Arc, + >::new()), + Vec::new(), + Vec::new(), + HashMap::from([(tool_name.clone(), vec![contributor.clone()])]), + ); + let tool_result = ResponseItem::FunctionCallOutput { + call_id: "call-1".to_string(), + output: FunctionCallOutputPayload::from_text("result".to_string()), + }; + let input = vec![ + ResponseItem::FunctionCall { + id: None, + name: tool_name.name, + namespace: tool_name.namespace, + arguments: "{}".to_string(), + call_id: "call-1".to_string(), + }, + tool_result.clone(), + ]; + + assert_eq!( + tool_result_contributors(&input, &router), + vec![UsagePromptContributor { + contributor, + source_estimated_tokens: estimate_response_item_tokens(&tool_result), + }] + ); + } + + fn usage_prompt_contributor( + kind: UsageContributorKind, + id: &str, + label: &str, + source_estimated_tokens: i64, + ) -> UsagePromptContributor { + UsagePromptContributor { + contributor: usage_contributor(kind, id, label), + source_estimated_tokens, + } + } + + fn usage_contributor(kind: UsageContributorKind, id: &str, label: &str) -> UsageContributor { + UsageContributor { + kind, + id: id.to_string(), + label: label.to_string(), + } + } +} diff --git a/codex-rs/protocol/src/protocol.rs b/codex-rs/protocol/src/protocol.rs index b0f34bae4d..c77de51620 100644 --- a/codex-rs/protocol/src/protocol.rs +++ b/codex-rs/protocol/src/protocol.rs @@ -1901,6 +1901,46 @@ pub struct TokenUsage { pub total_tokens: i64, } +#[derive(Debug, Clone, Copy, Deserialize, Serialize, PartialEq, Eq, Hash, JsonSchema, TS)] +#[serde(rename_all = "snake_case")] +#[ts(rename_all = "snake_case")] +pub enum UsageContributorKind { + Skill, + Subagent, + App, + McpServer, + Plugin, +} + +#[derive(Debug, Clone, Deserialize, Serialize, PartialEq, Eq, JsonSchema, TS)] +pub struct UsageContributor { + pub kind: UsageContributorKind, + pub id: String, + pub label: String, +} + +#[derive(Debug, Clone, Deserialize, Serialize, PartialEq, Eq, JsonSchema, TS)] +pub struct UsageAttributionContributor { + pub contributor: UsageContributor, + #[ts(type = "number")] + pub source_estimated_tokens: i64, + #[ts(type = "number")] + pub attributed_tokens: i64, +} + +#[derive(Debug, Clone, Deserialize, Serialize, PartialEq, Eq, JsonSchema, TS)] +pub struct UsageAttributionItem { + pub sample_id: String, + pub turn_id: String, + pub response_id: String, + #[ts(type = "number")] + pub occurred_at: i64, + pub token_usage: TokenUsage, + #[ts(type = "number")] + pub prompt_estimated_tokens: i64, + pub contributors: Vec, +} + #[derive(Debug, Clone, Deserialize, Serialize, PartialEq, Eq, JsonSchema, TS)] pub struct TokenUsageInfo { pub total_token_usage: TokenUsage, @@ -2688,6 +2728,7 @@ pub struct SessionMetaLine { pub enum RolloutItem { SessionMeta(SessionMetaLine), ResponseItem(ResponseItem), + UsageAttribution(UsageAttributionItem), Compacted(CompactedItem), TurnContext(TurnContextItem), EventMsg(EventMsg), diff --git a/codex-rs/rollout/src/list.rs b/codex-rs/rollout/src/list.rs index d4f0677898..a88c494f79 100644 --- a/codex-rs/rollout/src/list.rs +++ b/codex-rs/rollout/src/list.rs @@ -1135,6 +1135,9 @@ async fn read_head_summary(path: &Path, head_limit: usize) -> io::Result { // Not included in `head`; skip. } + RolloutItem::UsageAttribution(_) => { + // Not included in `head`; skip. + } RolloutItem::EventMsg(ev) => { if let Some(preview) = event_msg_preview(&ev) { if summary.preview.is_none() { @@ -1192,6 +1195,7 @@ pub async fn read_head_for_summary(path: &Path) -> io::Result {} } } diff --git a/codex-rs/rollout/src/metadata.rs b/codex-rs/rollout/src/metadata.rs index 2dd2df3a41..0889964979 100644 --- a/codex-rs/rollout/src/metadata.rs +++ b/codex-rs/rollout/src/metadata.rs @@ -72,6 +72,7 @@ pub fn builder_from_items( RolloutItem::ResponseItem(_) | RolloutItem::Compacted(_) | RolloutItem::TurnContext(_) + | RolloutItem::UsageAttribution(_) | RolloutItem::EventMsg(_) => None, }) && let Some(builder) = builder_from_session_meta(session_meta, rollout_path) { @@ -126,6 +127,7 @@ pub async fn extract_metadata_from_rollout( RolloutItem::ResponseItem(_) | RolloutItem::Compacted(_) | RolloutItem::TurnContext(_) + | RolloutItem::UsageAttribution(_) | RolloutItem::EventMsg(_) => None, }), parse_errors, diff --git a/codex-rs/rollout/src/policy.rs b/codex-rs/rollout/src/policy.rs index b6dbba98f8..d439288f1c 100644 --- a/codex-rs/rollout/src/policy.rs +++ b/codex-rs/rollout/src/policy.rs @@ -17,6 +17,7 @@ pub enum EventPersistenceMode { pub fn is_persisted_rollout_item(item: &RolloutItem, mode: EventPersistenceMode) -> bool { match item { RolloutItem::ResponseItem(item) => should_persist_response_item(item), + RolloutItem::UsageAttribution(_) => true, RolloutItem::EventMsg(ev) => should_persist_event_msg(ev, mode), // Persist Codex executive markers so we can analyze flows (e.g., compaction, API turns). RolloutItem::Compacted(_) | RolloutItem::TurnContext(_) | RolloutItem::SessionMeta(_) => { diff --git a/codex-rs/rollout/src/recorder.rs b/codex-rs/rollout/src/recorder.rs index 23b05b9def..d8eaca5fde 100644 --- a/codex-rs/rollout/src/recorder.rs +++ b/codex-rs/rollout/src/recorder.rs @@ -858,6 +858,9 @@ impl RolloutRecorder { RolloutItem::TurnContext(item) => { items.push(RolloutItem::TurnContext(item)); } + RolloutItem::UsageAttribution(item) => { + items.push(RolloutItem::UsageAttribution(item)); + } RolloutItem::EventMsg(_ev) => { items.push(RolloutItem::EventMsg(_ev)); } @@ -1735,6 +1738,7 @@ async fn resume_candidate_matches_cwd( RolloutItem::SessionMeta(_) | RolloutItem::ResponseItem(_) | RolloutItem::Compacted(_) + | RolloutItem::UsageAttribution(_) | RolloutItem::EventMsg(_) => None, }) { diff --git a/codex-rs/state/migrations/0035_usage_samples.sql b/codex-rs/state/migrations/0035_usage_samples.sql new file mode 100644 index 0000000000..eaae1f4fa7 --- /dev/null +++ b/codex-rs/state/migrations/0035_usage_samples.sql @@ -0,0 +1,29 @@ +CREATE TABLE usage_samples ( + sample_id TEXT PRIMARY KEY, + thread_id TEXT NOT NULL REFERENCES threads(id) ON DELETE CASCADE, + turn_id TEXT NOT NULL, + response_id TEXT NOT NULL, + occurred_at INTEGER NOT NULL, + input_tokens INTEGER NOT NULL, + cached_input_tokens INTEGER NOT NULL, + non_cached_input_tokens INTEGER NOT NULL, + output_tokens INTEGER NOT NULL, + reasoning_output_tokens INTEGER NOT NULL, + total_tokens INTEGER NOT NULL, + blended_tokens INTEGER NOT NULL, + prompt_estimated_tokens INTEGER NOT NULL +); + +CREATE TABLE usage_sample_contributors ( + sample_id TEXT NOT NULL REFERENCES usage_samples(sample_id) ON DELETE CASCADE, + kind TEXT NOT NULL, + contributor_id TEXT NOT NULL, + label TEXT NOT NULL, + source_estimated_tokens INTEGER NOT NULL, + attributed_tokens INTEGER NOT NULL, + PRIMARY KEY (sample_id, kind, contributor_id) +); + +CREATE INDEX idx_usage_samples_occurred_at ON usage_samples(occurred_at); +CREATE INDEX idx_usage_samples_thread_occurred_at ON usage_samples(thread_id, occurred_at); +CREATE INDEX idx_usage_sample_contributors_kind ON usage_sample_contributors(kind, contributor_id); diff --git a/codex-rs/state/src/extract.rs b/codex-rs/state/src/extract.rs index 2369addf23..976c6f8d81 100644 --- a/codex-rs/state/src/extract.rs +++ b/codex-rs/state/src/extract.rs @@ -22,6 +22,7 @@ pub fn apply_rollout_item( RolloutItem::TurnContext(turn_ctx) => apply_turn_context(metadata, turn_ctx), RolloutItem::EventMsg(event) => apply_event_msg(metadata, event), RolloutItem::ResponseItem(item) => apply_response_item(metadata, item), + RolloutItem::UsageAttribution(_) => {} RolloutItem::Compacted(_) => {} } if metadata.model_provider.is_empty() { @@ -36,9 +37,10 @@ pub fn rollout_item_affects_thread_metadata(item: &RolloutItem) -> bool { RolloutItem::EventMsg( EventMsg::TokenCount(_) | EventMsg::UserMessage(_) | EventMsg::ThreadGoalUpdated(_), ) => true, - RolloutItem::EventMsg(_) | RolloutItem::ResponseItem(_) | RolloutItem::Compacted(_) => { - false - } + RolloutItem::EventMsg(_) + | RolloutItem::ResponseItem(_) + | RolloutItem::UsageAttribution(_) + | RolloutItem::Compacted(_) => false, } } diff --git a/codex-rs/state/src/lib.rs b/codex-rs/state/src/lib.rs index f35b3420f6..bc8ba1467a 100644 --- a/codex-rs/state/src/lib.rs +++ b/codex-rs/state/src/lib.rs @@ -48,6 +48,11 @@ pub use model::ThreadGoalStatus; pub use model::ThreadMetadata; pub use model::ThreadMetadataBuilder; pub use model::ThreadsPage; +pub use model::UsageEntry; +pub use model::UsageHeadline; +pub use model::UsageRange; +pub use model::UsageReport; +pub use model::UsageSample; pub use runtime::GoalAccountingMode; pub use runtime::GoalAccountingOutcome; pub use runtime::GoalStore; diff --git a/codex-rs/state/src/model/mod.rs b/codex-rs/state/src/model/mod.rs index a431bc64c0..5fe3c11232 100644 --- a/codex-rs/state/src/model/mod.rs +++ b/codex-rs/state/src/model/mod.rs @@ -5,6 +5,7 @@ mod log; mod memories; mod thread_goal; mod thread_metadata; +mod usage; pub use agent_job::AgentJob; pub use agent_job::AgentJobCreateParams; @@ -34,6 +35,11 @@ pub use thread_metadata::SortKey; pub use thread_metadata::ThreadMetadata; pub use thread_metadata::ThreadMetadataBuilder; pub use thread_metadata::ThreadsPage; +pub use usage::UsageEntry; +pub use usage::UsageHeadline; +pub use usage::UsageRange; +pub use usage::UsageReport; +pub use usage::UsageSample; pub(crate) use agent_job::AgentJobItemRow; pub(crate) use agent_job::AgentJobRow; diff --git a/codex-rs/state/src/model/usage.rs b/codex-rs/state/src/model/usage.rs new file mode 100644 index 0000000000..23b9bee679 --- /dev/null +++ b/codex-rs/state/src/model/usage.rs @@ -0,0 +1,52 @@ +use codex_protocol::protocol::UsageAttributionItem; +use codex_protocol::protocol::UsageContributorKind; + +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum UsageRange { + Day, + Week, +} + +impl UsageRange { + pub(crate) fn seconds(self) -> i64 { + match self { + Self::Day => 24 * 60 * 60, + Self::Week => 7 * 24 * 60 * 60, + } + } +} + +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct UsageEntry { + pub kind: UsageContributorKind, + pub id: String, + pub label: String, + pub attributed_tokens: i64, + pub percent_of_usage: u8, +} + +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct UsageHeadline { + pub entry: UsageEntry, + pub note: Option, +} + +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct UsageReport { + pub range: UsageRange, + pub generated_at: i64, + pub tracked_from: Option, + pub total_tokens: i64, + pub headline: Option, + pub skills: Vec, + pub subagents: Vec, + pub apps: Vec, + pub mcp_servers: Vec, + pub plugins: Vec, +} + +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct UsageSample { + pub thread_id: codex_protocol::ThreadId, + pub attribution: UsageAttributionItem, +} diff --git a/codex-rs/state/src/runtime.rs b/codex-rs/state/src/runtime.rs index 8c9c6d0267..dd783696d0 100644 --- a/codex-rs/state/src/runtime.rs +++ b/codex-rs/state/src/runtime.rs @@ -65,6 +65,7 @@ mod remote_control; #[cfg(test)] mod test_support; mod threads; +mod usage; pub use goals::GoalAccountingMode; pub use goals::GoalAccountingOutcome; diff --git a/codex-rs/state/src/runtime/threads.rs b/codex-rs/state/src/runtime/threads.rs index c7030d7e44..28b8e56fd6 100644 --- a/codex-rs/state/src/runtime/threads.rs +++ b/codex-rs/state/src/runtime/threads.rs @@ -914,6 +914,15 @@ ON CONFLICT(thread_id, position) DO NOTHING { return Err(err); } + for item in items { + if let RolloutItem::UsageAttribution(attribution) = item { + self.record_usage_sample(&crate::UsageSample { + thread_id: builder.id, + attribution: attribution.clone(), + }) + .await?; + } + } Ok(()) } @@ -1036,6 +1045,7 @@ pub(super) fn extract_dynamic_tools(items: &[RolloutItem]) -> Option None, }) } @@ -1046,6 +1056,7 @@ pub(super) fn extract_memory_mode(items: &[RolloutItem]) -> Option { RolloutItem::ResponseItem(_) | RolloutItem::Compacted(_) | RolloutItem::TurnContext(_) + | RolloutItem::UsageAttribution(_) | RolloutItem::EventMsg(_) => None, }) } diff --git a/codex-rs/state/src/runtime/usage.rs b/codex-rs/state/src/runtime/usage.rs new file mode 100644 index 0000000000..dbedef6690 --- /dev/null +++ b/codex-rs/state/src/runtime/usage.rs @@ -0,0 +1,447 @@ +use super::*; +use crate::UsageEntry; +use crate::UsageHeadline; +use crate::UsageRange; +use crate::UsageReport; +use crate::UsageSample; +use codex_protocol::protocol::UsageContributorKind; + +impl StateRuntime { + pub async fn record_usage_sample(&self, sample: &UsageSample) -> anyhow::Result<()> { + let usage = &sample.attribution; + let token_usage = &usage.token_usage; + let mut tx = self.pool.begin().await?; + sqlx::query( + r#" +INSERT INTO usage_samples ( + sample_id, + thread_id, + turn_id, + response_id, + occurred_at, + input_tokens, + cached_input_tokens, + non_cached_input_tokens, + output_tokens, + reasoning_output_tokens, + total_tokens, + blended_tokens, + prompt_estimated_tokens +) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?) +ON CONFLICT(sample_id) DO UPDATE SET + thread_id = excluded.thread_id, + turn_id = excluded.turn_id, + response_id = excluded.response_id, + occurred_at = excluded.occurred_at, + input_tokens = excluded.input_tokens, + cached_input_tokens = excluded.cached_input_tokens, + non_cached_input_tokens = excluded.non_cached_input_tokens, + output_tokens = excluded.output_tokens, + reasoning_output_tokens = excluded.reasoning_output_tokens, + total_tokens = excluded.total_tokens, + blended_tokens = excluded.blended_tokens, + prompt_estimated_tokens = excluded.prompt_estimated_tokens + "#, + ) + .bind(usage.sample_id.as_str()) + .bind(sample.thread_id.to_string()) + .bind(usage.turn_id.as_str()) + .bind(usage.response_id.as_str()) + .bind(usage.occurred_at) + .bind(token_usage.input_tokens) + .bind(token_usage.cached_input_tokens) + .bind(token_usage.non_cached_input()) + .bind(token_usage.output_tokens) + .bind(token_usage.reasoning_output_tokens) + .bind(token_usage.total_tokens) + .bind(token_usage.blended_total()) + .bind(usage.prompt_estimated_tokens) + .execute(&mut *tx) + .await?; + sqlx::query("DELETE FROM usage_sample_contributors WHERE sample_id = ?") + .bind(usage.sample_id.as_str()) + .execute(&mut *tx) + .await?; + for contributor in &usage.contributors { + sqlx::query( + r#" +INSERT INTO usage_sample_contributors ( + sample_id, + kind, + contributor_id, + label, + source_estimated_tokens, + attributed_tokens +) VALUES (?, ?, ?, ?, ?, ?) + "#, + ) + .bind(usage.sample_id.as_str()) + .bind(usage_kind_key(contributor.contributor.kind)) + .bind(contributor.contributor.id.as_str()) + .bind(contributor.contributor.label.as_str()) + .bind(contributor.source_estimated_tokens) + .bind(contributor.attributed_tokens) + .execute(&mut *tx) + .await?; + } + tx.commit().await?; + Ok(()) + } + + pub async fn read_usage_report( + &self, + range: UsageRange, + now: i64, + ) -> anyhow::Result { + let since = now.saturating_sub(range.seconds()); + let total_tokens: i64 = sqlx::query_scalar( + "SELECT COALESCE(SUM(blended_tokens), 0) FROM usage_samples WHERE occurred_at >= ?", + ) + .bind(since) + .fetch_one(self.pool.as_ref()) + .await?; + let tracked_from: Option = + sqlx::query_scalar("SELECT MIN(occurred_at) FROM usage_samples") + .fetch_one(self.pool.as_ref()) + .await?; + let mut report = UsageReport { + range, + generated_at: now, + tracked_from, + total_tokens, + headline: None, + skills: self + .read_usage_contributors(since, UsageContributorKind::Skill, total_tokens) + .await?, + subagents: self.read_subagent_usage(since, total_tokens).await?, + apps: self + .read_usage_contributors(since, UsageContributorKind::App, total_tokens) + .await?, + mcp_servers: self + .read_usage_contributors(since, UsageContributorKind::McpServer, total_tokens) + .await?, + plugins: self + .read_usage_contributors(since, UsageContributorKind::Plugin, total_tokens) + .await?, + }; + report.headline = usage_headline(&report); + Ok(report) + } + + async fn read_usage_contributors( + &self, + since: i64, + kind: UsageContributorKind, + total_tokens: i64, + ) -> anyhow::Result> { + let rows = sqlx::query( + r#" +SELECT contributor_id, label, SUM(attributed_tokens) AS attributed_tokens +FROM usage_sample_contributors +JOIN usage_samples ON usage_samples.sample_id = usage_sample_contributors.sample_id +WHERE usage_samples.occurred_at >= ? + AND usage_sample_contributors.kind = ? +GROUP BY contributor_id, label +HAVING SUM(attributed_tokens) > 0 +ORDER BY attributed_tokens DESC, label ASC + "#, + ) + .bind(since) + .bind(usage_kind_key(kind)) + .fetch_all(self.pool.as_ref()) + .await?; + rows.into_iter() + .map(|row| { + let attributed_tokens = row.try_get("attributed_tokens")?; + Ok(UsageEntry { + kind, + id: row.try_get("contributor_id")?, + label: row.try_get("label")?, + attributed_tokens, + percent_of_usage: usage_percent(attributed_tokens, total_tokens), + }) + }) + .collect() + } + + async fn read_subagent_usage( + &self, + since: i64, + total_tokens: i64, + ) -> anyhow::Result> { + let rows = sqlx::query( + r#" +SELECT + COALESCE(NULLIF(threads.agent_role, ''), NULLIF(threads.agent_nickname, ''), 'subagent') AS label, + COALESCE(NULLIF(threads.agent_role, ''), NULLIF(threads.agent_nickname, ''), 'subagent') AS contributor_id, + SUM(usage_samples.blended_tokens) AS attributed_tokens +FROM usage_samples +JOIN threads ON threads.id = usage_samples.thread_id +WHERE usage_samples.occurred_at >= ? + AND threads.thread_source = 'subagent' +GROUP BY contributor_id, label +HAVING SUM(usage_samples.blended_tokens) > 0 +ORDER BY attributed_tokens DESC, label ASC + "#, + ) + .bind(since) + .fetch_all(self.pool.as_ref()) + .await?; + rows.into_iter() + .map(|row| { + let attributed_tokens = row.try_get("attributed_tokens")?; + Ok(UsageEntry { + kind: UsageContributorKind::Subagent, + id: row.try_get("contributor_id")?, + label: row.try_get("label")?, + attributed_tokens, + percent_of_usage: usage_percent(attributed_tokens, total_tokens), + }) + }) + .collect() + } +} + +fn usage_kind_key(kind: UsageContributorKind) -> &'static str { + match kind { + UsageContributorKind::Skill => "skill", + UsageContributorKind::Subagent => "subagent", + UsageContributorKind::App => "app", + UsageContributorKind::McpServer => "mcp_server", + UsageContributorKind::Plugin => "plugin", + } +} + +fn usage_percent(attributed_tokens: i64, total_tokens: i64) -> u8 { + if attributed_tokens <= 0 || total_tokens <= 0 { + return 0; + } + let rounded = attributed_tokens + .saturating_mul(/*rhs*/ 100) + .saturating_add(total_tokens / 2) + / total_tokens; + u8::try_from(rounded.max(/*other*/ 1).min(i64::from(u8::MAX))).unwrap_or(u8::MAX) +} + +fn usage_headline(report: &UsageReport) -> Option { + let entry = report + .skills + .iter() + .chain(report.subagents.iter()) + .chain(report.apps.iter()) + .chain(report.mcp_servers.iter()) + .chain(report.plugins.iter()) + .max_by(|left, right| { + left.attributed_tokens + .cmp(&right.attributed_tokens) + .then_with(|| right.label.cmp(&left.label)) + })? + .clone(); + let note = matches!( + entry.kind, + UsageContributorKind::App | UsageContributorKind::McpServer + ) + .then(|| { + "Tool results stay in context until compaction; compact or disable sources you do not need." + .to_string() + }); + Some(UsageHeadline { entry, note }) +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::runtime::test_support::test_thread_metadata; + use crate::runtime::test_support::unique_temp_dir; + use codex_protocol::ThreadId; + use codex_protocol::protocol::ThreadSource; + use codex_protocol::protocol::TokenUsage; + use codex_protocol::protocol::UsageAttributionContributor; + use codex_protocol::protocol::UsageAttributionItem; + use codex_protocol::protocol::UsageContributor; + use pretty_assertions::assert_eq; + + #[tokio::test] + async fn usage_report_groups_forward_only_samples_by_range() { + let codex_home = unique_temp_dir(); + let runtime = StateRuntime::init(codex_home.clone(), "test-provider".to_string()) + .await + .expect("state db should initialize"); + let user_thread_id = + ThreadId::from_string("00000000-0000-0000-0000-000000000901").expect("valid thread id"); + let subagent_thread_id = + ThreadId::from_string("00000000-0000-0000-0000-000000000902").expect("valid thread id"); + let now = 1_700_000_200; + + runtime + .upsert_thread(&test_thread_metadata( + &codex_home, + user_thread_id, + codex_home.clone(), + )) + .await + .expect("user thread insert should succeed"); + let mut subagent_metadata = + test_thread_metadata(&codex_home, subagent_thread_id, codex_home.clone()); + subagent_metadata.thread_source = Some(ThreadSource::Subagent); + subagent_metadata.agent_role = Some("code-review".to_string()); + runtime + .upsert_thread(&subagent_metadata) + .await + .expect("subagent thread insert should succeed"); + + runtime + .record_usage_sample(&usage_sample( + user_thread_id, + "recent-user", + /*occurred_at*/ now - 100, + TokenUsage { + input_tokens: 100, + cached_input_tokens: 20, + output_tokens: 40, + reasoning_output_tokens: 0, + total_tokens: 140, + }, + vec![ + contributor( + UsageContributorKind::Skill, + "/skills/tmux", + "tmux", + /*attributed_tokens*/ 50, + ), + contributor( + UsageContributorKind::App, + "slack", + "Slack", + /*attributed_tokens*/ 70, + ), + ], + )) + .await + .expect("recent usage sample should persist"); + runtime + .record_usage_sample(&usage_sample( + subagent_thread_id, + "recent-subagent", + /*occurred_at*/ now - 50, + TokenUsage { + input_tokens: 30, + cached_input_tokens: 0, + output_tokens: 10, + reasoning_output_tokens: 0, + total_tokens: 40, + }, + Vec::new(), + )) + .await + .expect("subagent usage sample should persist"); + runtime + .record_usage_sample(&usage_sample( + user_thread_id, + "old-user", + /*occurred_at*/ now - UsageRange::Day.seconds() - 1, + TokenUsage { + input_tokens: 10, + cached_input_tokens: 0, + output_tokens: 0, + reasoning_output_tokens: 0, + total_tokens: 10, + }, + vec![contributor( + UsageContributorKind::McpServer, + "old-mcp", + "old-mcp", + /*attributed_tokens*/ 10, + )], + )) + .await + .expect("old usage sample should persist"); + + assert_eq!( + runtime + .read_usage_report(UsageRange::Day, now) + .await + .expect("usage report should load"), + UsageReport { + range: UsageRange::Day, + generated_at: now, + tracked_from: Some(now - UsageRange::Day.seconds() - 1), + total_tokens: 160, + headline: Some(UsageHeadline { + entry: UsageEntry { + kind: UsageContributorKind::App, + id: "slack".to_string(), + label: "Slack".to_string(), + attributed_tokens: 70, + percent_of_usage: 44, + }, + note: Some( + "Tool results stay in context until compaction; compact or disable sources you do not need." + .to_string(), + ), + }), + skills: vec![UsageEntry { + kind: UsageContributorKind::Skill, + id: "/skills/tmux".to_string(), + label: "tmux".to_string(), + attributed_tokens: 50, + percent_of_usage: 31, + }], + subagents: vec![UsageEntry { + kind: UsageContributorKind::Subagent, + id: "code-review".to_string(), + label: "code-review".to_string(), + attributed_tokens: 40, + percent_of_usage: 25, + }], + apps: vec![UsageEntry { + kind: UsageContributorKind::App, + id: "slack".to_string(), + label: "Slack".to_string(), + attributed_tokens: 70, + percent_of_usage: 44, + }], + mcp_servers: Vec::new(), + plugins: Vec::new(), + } + ); + } + + fn usage_sample( + thread_id: ThreadId, + sample_id: &str, + occurred_at: i64, + token_usage: TokenUsage, + contributors: Vec, + ) -> UsageSample { + UsageSample { + thread_id, + attribution: UsageAttributionItem { + sample_id: sample_id.to_string(), + turn_id: format!("{sample_id}-turn"), + response_id: format!("{sample_id}-response"), + occurred_at, + token_usage, + prompt_estimated_tokens: 100, + contributors, + }, + } + } + + fn contributor( + kind: UsageContributorKind, + id: &str, + label: &str, + attributed_tokens: i64, + ) -> UsageAttributionContributor { + UsageAttributionContributor { + contributor: UsageContributor { + kind, + id: id.to_string(), + label: label.to_string(), + }, + source_estimated_tokens: attributed_tokens, + attributed_tokens, + } + } +} diff --git a/codex-rs/thread-store/src/thread_metadata_sync.rs b/codex-rs/thread-store/src/thread_metadata_sync.rs index edc78a36c2..a614c1a3a4 100644 --- a/codex-rs/thread-store/src/thread_metadata_sync.rs +++ b/codex-rs/thread-store/src/thread_metadata_sync.rs @@ -278,6 +278,7 @@ impl ThreadMetadataSync { RolloutItem::SessionMeta(_) | RolloutItem::EventMsg(_) | RolloutItem::ResponseItem(_) + | RolloutItem::UsageAttribution(_) | RolloutItem::Compacted(_) => {} } } diff --git a/codex-rs/tui/src/app/background_requests.rs b/codex-rs/tui/src/app/background_requests.rs index dcf1f672be..072614098e 100644 --- a/codex-rs/tui/src/app/background_requests.rs +++ b/codex-rs/tui/src/app/background_requests.rs @@ -15,6 +15,9 @@ use codex_app_server_protocol::MarketplaceRemoveParams; use codex_app_server_protocol::MarketplaceRemoveResponse; use codex_app_server_protocol::MarketplaceUpgradeParams; use codex_app_server_protocol::MarketplaceUpgradeResponse; +use codex_app_server_protocol::UsageRange; +use codex_app_server_protocol::UsageReadParams; +use codex_app_server_protocol::UsageReadResponse; use codex_app_server_protocol::RequestId; @@ -127,6 +130,26 @@ impl App { }); } + pub(super) fn fetch_usage( + &mut self, + app_server: &AppServerSession, + request_id: u64, + range: UsageRange, + ) { + let request_handle = app_server.request_handle(); + let app_event_tx = self.app_event_tx.clone(); + tokio::spawn(async move { + let result = fetch_usage(request_handle, range) + .await + .map_err(|err| err.to_string()); + app_event_tx.send(AppEvent::UsageLoaded { + request_id, + range, + result, + }); + }); + } + pub(super) fn fetch_hooks_list(&mut self, app_server: &AppServerSession, cwd: PathBuf) { let request_handle = app_server.request_handle(); let app_event_tx = self.app_event_tx.clone(); @@ -577,6 +600,20 @@ impl App { } } +async fn fetch_usage( + request_handle: AppServerRequestHandle, + range: UsageRange, +) -> Result { + let request_id = RequestId::String(format!("usage-read-{}", uuid::Uuid::new_v4())); + request_handle + .request_typed(ClientRequest::UsageRead { + request_id, + params: UsageReadParams { range }, + }) + .await + .map_err(Into::into) +} + pub(super) async fn fetch_all_mcp_server_statuses( request_handle: AppServerRequestHandle, detail: McpServerStatusDetail, diff --git a/codex-rs/tui/src/app/event_dispatch.rs b/codex-rs/tui/src/app/event_dispatch.rs index bdd7b46917..8bf32db265 100644 --- a/codex-rs/tui/src/app/event_dispatch.rs +++ b/codex-rs/tui/src/app/event_dispatch.rs @@ -436,6 +436,19 @@ impl App { AppEvent::FetchPluginsList { cwd } => { self.fetch_plugins_list(app_server, cwd); } + AppEvent::OpenUsage { range } => { + self.chat_widget.open_usage(range); + } + AppEvent::FetchUsage { request_id, range } => { + self.fetch_usage(app_server, request_id, range); + } + AppEvent::UsageLoaded { + request_id, + range, + result, + } => { + self.chat_widget.on_usage_loaded(request_id, range, result); + } AppEvent::FetchHooksList { cwd } => { self.fetch_hooks_list(app_server, cwd); } diff --git a/codex-rs/tui/src/app_event.rs b/codex-rs/tui/src/app_event.rs index a85ee5a70c..7fc95be534 100644 --- a/codex-rs/tui/src/app_event.rs +++ b/codex-rs/tui/src/app_event.rs @@ -371,6 +371,21 @@ pub(crate) enum AppEvent { cwd: PathBuf, }, + OpenUsage { + range: codex_app_server_protocol::UsageRange, + }, + + FetchUsage { + request_id: u64, + range: codex_app_server_protocol::UsageRange, + }, + + UsageLoaded { + request_id: u64, + range: codex_app_server_protocol::UsageRange, + result: Result, + }, + /// Fetch lifecycle hook inventory for the provided working directory. FetchHooksList { cwd: PathBuf, diff --git a/codex-rs/tui/src/bottom_pane/mod.rs b/codex-rs/tui/src/bottom_pane/mod.rs index 65770b61ae..40302ab259 100644 --- a/codex-rs/tui/src/bottom_pane/mod.rs +++ b/codex-rs/tui/src/bottom_pane/mod.rs @@ -58,6 +58,7 @@ mod status_line_setup; mod status_line_style; mod status_surface_preview; mod title_setup; +mod usage_view; pub(crate) use action_required_title::ACTION_REQUIRED_PREVIEW_PREFIX; pub(crate) use action_required_title::build_action_required_title_text; pub(crate) use app_link_view::AppLinkElicitationTarget; @@ -71,6 +72,8 @@ pub(crate) use mcp_server_elicitation::McpServerElicitationFormRequest; pub(crate) use mcp_server_elicitation::McpServerElicitationOverlay; pub(crate) use request_user_input::RequestUserInputOverlay; pub(crate) use status_line_style::status_line_from_segments; +pub(crate) use usage_view::USAGE_VIEW_ID; +pub(crate) use usage_view::UsageView; mod bottom_pane_view; #[derive(Clone, Debug, PartialEq, Eq)] diff --git a/codex-rs/tui/src/bottom_pane/usage_view.rs b/codex-rs/tui/src/bottom_pane/usage_view.rs new file mode 100644 index 0000000000..7cd9f9d536 --- /dev/null +++ b/codex-rs/tui/src/bottom_pane/usage_view.rs @@ -0,0 +1,146 @@ +use crate::app_event::AppEvent; +use crate::app_event_sender::AppEventSender; +use crate::bottom_pane::BottomPaneView; +use crate::render::renderable::Renderable; +use codex_app_server_protocol::UsageEntry; +use codex_app_server_protocol::UsageRange; +use codex_app_server_protocol::UsageReadResponse; +use crossterm::event::KeyCode; +use crossterm::event::KeyEvent; +use crossterm::event::KeyEventKind; +use ratatui::buffer::Buffer; +use ratatui::layout::Rect; +use ratatui::style::Stylize; +use ratatui::text::Line; +use ratatui::widgets::Paragraph; +use ratatui::widgets::WidgetRef; + +pub(crate) const USAGE_VIEW_ID: &str = "usage"; + +pub(crate) struct UsageView { + lines: Vec>, + range: UsageRange, + app_event_tx: AppEventSender, +} + +impl UsageView { + pub(crate) fn loading(range: UsageRange, app_event_tx: AppEventSender) -> Self { + Self { + lines: vec![Line::from("Usage".bold()), Line::from("Loading...".dim())], + range, + app_event_tx, + } + } + + pub(crate) fn loaded(response: UsageReadResponse, app_event_tx: AppEventSender) -> Self { + let range = response.report.range; + let report = response.report; + let mut lines = vec![Line::from("Usage".bold())]; + if let Some(headline) = report.headline.as_ref() { + lines.push(Line::from(format!( + "{}% of your usage came from {} \"{}\"", + headline.entry.percent_of_usage, + contributor_kind_label(headline.entry.kind), + headline.entry.label + ))); + if let Some(note) = headline.note.as_ref() { + lines.push(Line::from(note.clone().dim())); + } + lines.push(Line::from("")); + } + if report.total_tokens == 0 { + lines.push(Line::from("No tracked usage in this range yet.".dim())); + } else { + push_section(&mut lines, "Skills", &report.skills); + push_section(&mut lines, "Subagents", &report.subagents); + push_section(&mut lines, "Apps", &report.apps); + push_section(&mut lines, "MCP servers", &report.mcp_servers); + push_section(&mut lines, "Plugins", &report.plugins); + } + if lines.last().is_none_or(|line| line.width() != 0) { + lines.push(Line::from("")); + } + lines.push(Line::from(range_footer(range).dim())); + Self { + lines, + range, + app_event_tx, + } + } + + pub(crate) fn error(range: UsageRange, err: String, app_event_tx: AppEventSender) -> Self { + Self { + lines: vec![ + Line::from("Usage".bold()), + Line::from(format!("Failed to load usage: {err}").red()), + Line::from(""), + Line::from(range_footer(range).dim()), + ], + range, + app_event_tx, + } + } +} + +impl BottomPaneView for UsageView { + fn handle_key_event(&mut self, key_event: KeyEvent) { + if key_event.kind == KeyEventKind::Release { + return; + } + match key_event.code { + KeyCode::Char('d') if self.range != UsageRange::Day => { + self.app_event_tx.send(AppEvent::OpenUsage { + range: UsageRange::Day, + }); + } + KeyCode::Char('w') if self.range != UsageRange::Week => { + self.app_event_tx.send(AppEvent::OpenUsage { + range: UsageRange::Week, + }); + } + _ => {} + } + } + + fn view_id(&self) -> Option<&'static str> { + Some(USAGE_VIEW_ID) + } +} + +impl Renderable for UsageView { + fn render(&self, area: Rect, buf: &mut Buffer) { + Paragraph::new(self.lines.clone()).render_ref(area, buf); + } + + fn desired_height(&self, _width: u16) -> u16 { + u16::try_from(self.lines.len()).unwrap_or(u16::MAX) + } +} + +fn push_section(lines: &mut Vec>, label: &str, entries: &[UsageEntry]) { + if entries.is_empty() { + return; + } + lines.push(Line::from(label.to_string().bold())); + for entry in entries { + lines.push(Line::from(format!( + "{:<24} {:>3}%", + entry.label, entry.percent_of_usage + ))); + } + lines.push(Line::from("")); +} + +fn contributor_kind_label(kind: codex_app_server_protocol::UsageContributorKind) -> &'static str { + match kind { + codex_app_server_protocol::UsageContributorKind::Skill => "skill", + codex_app_server_protocol::UsageContributorKind::Subagent => "subagent", + codex_app_server_protocol::UsageContributorKind::App => "app", + codex_app_server_protocol::UsageContributorKind::McpServer => "MCP server", + codex_app_server_protocol::UsageContributorKind::Plugin => "plugin", + } +} + +fn range_footer(_range: UsageRange) -> &'static str { + "d to day | w to week" +} diff --git a/codex-rs/tui/src/chatwidget.rs b/codex-rs/tui/src/chatwidget.rs index 0ba974d88e..3b44c966f1 100644 --- a/codex-rs/tui/src/chatwidget.rs +++ b/codex-rs/tui/src/chatwidget.rs @@ -361,6 +361,7 @@ use self::skills::collect_tool_mentions; use self::skills::find_app_mentions; use self::skills::find_skill_mentions_with_tool_mentions; mod plugins; +mod usage; use self::plugins::PluginInstallAuthFlowState; use self::plugins::PluginListFetchState; use self::plugins::PluginsCacheState; @@ -537,6 +538,8 @@ pub(crate) struct ChatWidget { rate_limit_snapshots_by_limit_id: BTreeMap, refreshing_status_outputs: Vec<(u64, StatusHistoryHandle)>, next_status_refresh_request_id: u64, + next_usage_request_id: u64, + active_usage_request_id: Option, plan_type: Option, codex_rate_limit_reached_type: Option, rate_limit_warnings: RateLimitWarningState, diff --git a/codex-rs/tui/src/chatwidget/constructor.rs b/codex-rs/tui/src/chatwidget/constructor.rs index ef90c85d4c..dd4ba02fa3 100644 --- a/codex-rs/tui/src/chatwidget/constructor.rs +++ b/codex-rs/tui/src/chatwidget/constructor.rs @@ -123,6 +123,8 @@ impl ChatWidget { rate_limit_snapshots_by_limit_id: BTreeMap::new(), refreshing_status_outputs: Vec::new(), next_status_refresh_request_id: 0, + next_usage_request_id: 0, + active_usage_request_id: None, plan_type: initial_plan_type, codex_rate_limit_reached_type: None, rate_limit_warnings: RateLimitWarningState::default(), diff --git a/codex-rs/tui/src/chatwidget/slash_dispatch.rs b/codex-rs/tui/src/chatwidget/slash_dispatch.rs index aa9e5112f8..a956c0c29a 100644 --- a/codex-rs/tui/src/chatwidget/slash_dispatch.rs +++ b/codex-rs/tui/src/chatwidget/slash_dispatch.rs @@ -382,6 +382,9 @@ impl ChatWidget { ); } } + SlashCommand::Usage => { + self.open_usage(codex_app_server_protocol::UsageRange::Day); + } SlashCommand::Ide => { self.handle_ide_command(); } @@ -929,6 +932,7 @@ impl ChatWidget { match cmd { SlashCommand::Ide | SlashCommand::Status + | SlashCommand::Usage | SlashCommand::DebugConfig | SlashCommand::Ps | SlashCommand::Stop diff --git a/codex-rs/tui/src/chatwidget/snapshots/codex_tui__chatwidget__tests__usage_popup.snap b/codex-rs/tui/src/chatwidget/snapshots/codex_tui__chatwidget__tests__usage_popup.snap new file mode 100644 index 0000000000..18a741f4cf --- /dev/null +++ b/codex-rs/tui/src/chatwidget/snapshots/codex_tui__chatwidget__tests__usage_popup.snap @@ -0,0 +1,20 @@ +--- +source: tui/src/chatwidget/tests/popups_and_settings.rs +expression: "render_bottom_popup(&chat, 100)" +--- +Usage +11% of your usage came from app "testmcp" +Tool results stay in context until compaction; compact or disable sources you do not need. + +Skills +/tmux 8% +/babysit 6% + +Subagents +babysit 13% +code-review 9% + +Apps +testmcp 11% + +d to day | w to week diff --git a/codex-rs/tui/src/chatwidget/tests.rs b/codex-rs/tui/src/chatwidget/tests.rs index 5d27e0ddc5..d5536eac58 100644 --- a/codex-rs/tui/src/chatwidget/tests.rs +++ b/codex-rs/tui/src/chatwidget/tests.rs @@ -116,6 +116,12 @@ pub(super) use codex_app_server_protocol::TurnCompletedNotification; pub(super) use codex_app_server_protocol::TurnError as AppServerTurnError; pub(super) use codex_app_server_protocol::TurnStartedNotification; pub(super) use codex_app_server_protocol::TurnStatus as AppServerTurnStatus; +pub(super) use codex_app_server_protocol::UsageContributorKind; +pub(super) use codex_app_server_protocol::UsageEntry; +pub(super) use codex_app_server_protocol::UsageHeadline; +pub(super) use codex_app_server_protocol::UsageRange; +pub(super) use codex_app_server_protocol::UsageReadResponse; +pub(super) use codex_app_server_protocol::UsageReport; pub(super) use codex_app_server_protocol::UserInput; pub(super) use codex_app_server_protocol::UserInput as AppServerUserInput; pub(super) use codex_app_server_protocol::WarningNotification; diff --git a/codex-rs/tui/src/chatwidget/tests/popups_and_settings.rs b/codex-rs/tui/src/chatwidget/tests/popups_and_settings.rs index 09d911ad16..f9c523d562 100644 --- a/codex-rs/tui/src/chatwidget/tests/popups_and_settings.rs +++ b/codex-rs/tui/src/chatwidget/tests/popups_and_settings.rs @@ -111,6 +111,82 @@ async fn plugins_popup_loading_state_snapshot() { assert_chatwidget_snapshot!("plugins_popup_loading_state", popup); } +#[tokio::test] +async fn usage_popup_snapshot() { + let (mut chat, mut rx, _op_rx) = make_chatwidget_manual(/*model_override*/ None).await; + + chat.open_usage(UsageRange::Day); + assert_matches!( + rx.try_recv(), + Ok(AppEvent::FetchUsage { + request_id: 0, + range: UsageRange::Day, + }) + ); + chat.on_usage_loaded( + /*request_id*/ 0, + UsageRange::Day, + Ok(UsageReadResponse { + report: UsageReport { + range: UsageRange::Day, + generated_at: 1_700_000_000, + tracked_from: Some(1_699_999_000), + total_tokens: 100, + headline: Some(UsageHeadline { + entry: usage_entry( + UsageContributorKind::App, + "testmcp", + "testmcp", + /*percent_of_usage*/ 11, + ), + note: Some( + "Tool results stay in context until compaction; compact or disable sources you do not need." + .to_string(), + ), + }), + skills: vec![ + usage_entry( + UsageContributorKind::Skill, + "/skills/tmux", + "/tmux", + /*percent_of_usage*/ 8, + ), + usage_entry( + UsageContributorKind::Skill, + "/skills/babysit", + "/babysit", + /*percent_of_usage*/ 6, + ), + ], + subagents: vec![ + usage_entry( + UsageContributorKind::Subagent, + "babysit", + "babysit", + /*percent_of_usage*/ 13, + ), + usage_entry( + UsageContributorKind::Subagent, + "code-review", + "code-review", + /*percent_of_usage*/ 9, + ), + ], + apps: vec![usage_entry( + UsageContributorKind::App, + "testmcp", + "testmcp", + /*percent_of_usage*/ 11, + )], + mcp_servers: Vec::new(), + plugins: Vec::new(), + }, + }), + ); + + assert_chatwidget_snapshot!("usage_popup", render_bottom_popup(&chat, /*width*/ 100)); +} + #[tokio::test] async fn marketplace_upgrade_loading_popup_snapshot() { let (mut chat, _rx, _op_rx) = make_chatwidget_manual(/*model_override*/ None).await; @@ -131,6 +207,21 @@ async fn marketplace_upgrade_loading_popup_snapshot() { ); } +fn usage_entry( + kind: UsageContributorKind, + id: &str, + label: &str, + percent_of_usage: u8, +) -> UsageEntry { + UsageEntry { + kind, + id: id.to_string(), + label: label.to_string(), + attributed_tokens: i64::from(percent_of_usage), + percent_of_usage, + } +} + #[tokio::test] async fn marketplace_upgrade_failure_includes_backend_messages_snapshot() { let (mut chat, mut rx, _op_rx) = make_chatwidget_manual(/*model_override*/ None).await; diff --git a/codex-rs/tui/src/chatwidget/usage.rs b/codex-rs/tui/src/chatwidget/usage.rs new file mode 100644 index 0000000000..98b5ad416b --- /dev/null +++ b/codex-rs/tui/src/chatwidget/usage.rs @@ -0,0 +1,41 @@ +use super::ChatWidget; +use crate::app_event::AppEvent; +use crate::bottom_pane::USAGE_VIEW_ID; +use crate::bottom_pane::UsageView; +use codex_app_server_protocol::UsageRange; +use codex_app_server_protocol::UsageReadResponse; + +impl ChatWidget { + pub(crate) fn open_usage(&mut self, range: UsageRange) { + let request_id = self.next_usage_request_id; + self.next_usage_request_id = self.next_usage_request_id.saturating_add(/*rhs*/ 1); + self.active_usage_request_id = Some(request_id); + self.bottom_pane.dismiss_active_view_if_id(USAGE_VIEW_ID); + self.bottom_pane.show_view(Box::new(UsageView::loading( + range, + self.app_event_tx.clone(), + ))); + self.app_event_tx + .send(AppEvent::FetchUsage { request_id, range }); + self.request_redraw(); + } + + pub(crate) fn on_usage_loaded( + &mut self, + request_id: u64, + range: UsageRange, + result: Result, + ) { + if self.active_usage_request_id != Some(request_id) + || !self.bottom_pane.dismiss_active_view_if_id(USAGE_VIEW_ID) + { + return; + } + let view = match result { + Ok(response) => UsageView::loaded(response, self.app_event_tx.clone()), + Err(err) => UsageView::error(range, err, self.app_event_tx.clone()), + }; + self.bottom_pane.show_view(Box::new(view)); + self.request_redraw(); + } +} diff --git a/codex-rs/tui/src/slash_command.rs b/codex-rs/tui/src/slash_command.rs index 4a43409296..caf09d25a5 100644 --- a/codex-rs/tui/src/slash_command.rs +++ b/codex-rs/tui/src/slash_command.rs @@ -44,6 +44,7 @@ pub enum SlashCommand { Diff, Mention, Status, + Usage, DebugConfig, Title, Statusline, @@ -96,6 +97,9 @@ impl SlashCommand { SlashCommand::Skills => "use skills to improve how Codex performs specific tasks", SlashCommand::Hooks => "view and manage lifecycle hooks", SlashCommand::Status => "show current session configuration and token usage", + SlashCommand::Usage => { + "show local token usage by skills, subagents, apps, MCP servers, and plugins" + } SlashCommand::DebugConfig => "show config layers and requirement sources for debugging", SlashCommand::Title => "configure which items appear in the terminal title", SlashCommand::Statusline => "configure which items appear in the status line", @@ -172,6 +176,7 @@ impl SlashCommand { | SlashCommand::Diff | SlashCommand::Mention | SlashCommand::Status + | SlashCommand::Usage | SlashCommand::Ide ) } @@ -207,6 +212,7 @@ impl SlashCommand { | SlashCommand::Skills | SlashCommand::Hooks | SlashCommand::Status + | SlashCommand::Usage | SlashCommand::DebugConfig | SlashCommand::Ps | SlashCommand::Stop