@sixb/core 0.1.6 → 0.1.7

package/src/storage/in-memory/index.ts

@@ -2,6 +2,7 @@ import { AsyncLocalStorage } from "node:async_hooks"
 import { InMemoryActionRunStorage } from "../action-runs"
 import { type AgentStorage, InMemoryAgentStorage } from "../agents"
 import { type AiAccountingAttribution, type AiCostStorage, InMemoryAiCostStorage } from "../ai-cost"
+import { type AiLimitStorage, InMemoryAiLimitStorage } from "../ai-limits"
 import { type AiUsageStorage, InMemoryAiUsageStorage } from "../ai-usage"
 import { type AuthStorage, InMemoryAuthStorage } from "../auth"
 import {
@@ -70,6 +71,31 @@ export class InMemoryStorage implements Storage {
   private readonly aiUsageStorage = new InMemoryAiUsageStorage(this.executionStorage)
   private readonly aiCostStorage = new InMemoryAiCostStorage(this.aiUsageStorage, {
     resolve: (input) => this.resolveAiAccountingAttribution(input),
+    resolveGroup: async (input) => {
+      const attribution = await this.resolveAiAccountingAttribution(input)
+      if (attribution?.kind !== "subagent") return { executionId: input.executionId, attribution }
+      const parent = await this.agentStorage.runs.getById({
+        projectId: input.projectId,
+        id: attribution.parentRunId,
+      })
+      if (!parent) throw new Error("[Sixb] Subagent accounting parent is missing.")
+      return {
+        executionId: parent.executionId,
+        attribution: await this.resolveAiAccountingAttribution({
+          ...input,
+          executionId: parent.executionId,
+        }),
+      }
+    },
+  })
+  private readonly aiLimitStorage = new InMemoryAiLimitStorage({
+    resolveExecution: ({ projectId, executionId }) =>
+      this.executionStorage.getById({ projectId, id: executionId }),
+    resolveUsageRecord: ({ projectId, usageRecordId }) =>
+      this.aiUsageStorage.getRecord(projectId, usageRecordId),
+    resolveCostRecord: ({ projectId, usageRecordId }) =>
+      this.aiCostStorage.getRecord(projectId, usageRecordId),
+    listUsageRecords: ({ projectId }) => this.aiUsageStorage.listRecords(projectId),
   })
   private readonly actionRunStorage = new InMemoryActionRunStorage(this.executionStorage)
   private readonly syncRunStorage = new InMemorySyncRunStorage(this.executionStorage)
@@ -88,6 +114,7 @@ export class InMemoryStorage implements Storage {
   readonly agents: AgentStorage
   readonly aiUsage: AiUsageStorage
   readonly aiCosts: AiCostStorage
+  readonly aiLimits: AiLimitStorage
   readonly actionRuns: InMemoryActionRunStorage
   readonly syncRuns: SyncRunStorage
   readonly pipelineRuns: PipelineRunStorage
@@ -114,6 +141,7 @@ export class InMemoryStorage implements Storage {
     this.agents = createAgentOperationScope(this.agentStorage, scope)
     this.aiUsage = createOperationScopedFacade(this.aiUsageStorage, scope)
     this.aiCosts = createOperationScopedFacade(this.aiCostStorage, scope)
+    this.aiLimits = createOperationScopedFacade(this.aiLimitStorage, scope)
     this.actionRuns = createOperationScopedFacade(this.actionRunStorage, scope)
     this.syncRuns = createOperationScopedFacade(this.syncRunStorage, scope)
     this.pipelineRuns = createOperationScopedFacade(this.pipelineRunStorage, scope)
@@ -154,14 +182,20 @@ export class InMemoryStorage implements Storage {
       projectId: input.projectId,
       id: execution.executor.runId,
     })
-    if (directRun) {
+    if (directRun?.kind === "conversation") {
       return {
         kind: "agent",
-        agentId: directRun.agentId,
         agentRunId: directRun.id,
         threadId: directRun.threadId,
       }
     }
+    if (directRun?.kind === "subagent") {
+      return {
+        kind: "subagent",
+        subagentRunId: directRun.id,
+        parentRunId: directRun.parentRunId,
+      }
+    }
     const agentRun = await this.workflowRunStorage.agentNodes.getByNodeRunId({
       projectId: input.projectId,
       nodeRunId: execution.executor.runId,
@@ -174,7 +208,7 @@ export class InMemoryStorage implements Storage {
     if (!node) return undefined
     return {
       kind: "workflowAgent",
-      agentId: agentRun.agentId,
+      agentStepId: node.nodeId,
       nodeRunId: agentRun.nodeRunId,
       workflowId: node.workflowId,
       workflowRunId: node.workflowRunId,
@@ -224,7 +258,13 @@ export class InMemoryStorage implements Storage {
       this.materializationLifecycles.set(transactionToken, materializationLifecycle)
 
       try {
-        const result = await this.transactionScope.run(transactionToken, async () => run(tx))
+        const result = await this.transactionScope.run(transactionToken, async () => {
+          try {
+            return await run(tx)
+          } finally {
+            await this.aiLimitStorage.settleOperations()
+          }
+        })
         materializationLifecycle.assertCommittable()
         return result
       } catch (error) {
@@ -297,6 +337,7 @@ export class InMemoryStorage implements Storage {
       agents: this.agentStorage,
       aiUsage: this.aiUsageStorage,
       aiCosts: this.aiCostStorage,
+      aiLimits: this.aiLimitStorage,
       actionRuns: this.actionRunStorage,
       syncRuns: this.syncRunStorage,
       pipelineRuns: this.pipelineRunStorage,
@@ -324,6 +365,7 @@ export class InMemoryStorage implements Storage {
       agents: this.agentStorage.snapshot(),
       aiUsage: this.aiUsageStorage.snapshot(),
       aiCosts: this.aiCostStorage.snapshot(),
+      aiLimits: this.aiLimitStorage.snapshot(),
       actionRuns: this.actionRunStorage.snapshot(),
       syncRuns: this.syncRunStorage.snapshot(),
       pipelineRuns: this.pipelineRunStorage.snapshot(),
@@ -346,6 +388,7 @@ export class InMemoryStorage implements Storage {
     this.agentStorage.restore(snapshot.agents)
     this.aiUsageStorage.restore(snapshot.aiUsage)
     this.aiCostStorage.restore(snapshot.aiCosts)
+    this.aiLimitStorage.restore(snapshot.aiLimits)
     this.actionRunStorage.restore(snapshot.actionRuns)
     this.syncRunStorage.restore(snapshot.syncRuns)
     this.pipelineRunStorage.restore(snapshot.pipelineRuns)
@@ -368,6 +411,7 @@ export interface InMemoryStorageSnapshot {
   readonly agents: ReturnType<InMemoryAgentStorage["snapshot"]>
   readonly aiUsage: ReturnType<InMemoryAiUsageStorage["snapshot"]>
   readonly aiCosts: ReturnType<InMemoryAiCostStorage["snapshot"]>
+  readonly aiLimits: ReturnType<InMemoryAiLimitStorage["snapshot"]>
   readonly actionRuns: ReturnType<InMemoryActionRunStorage["snapshot"]>
   readonly syncRuns: ReturnType<InMemorySyncRunStorage["snapshot"]>
   readonly pipelineRuns: ReturnType<InMemoryPipelineRunStorage["snapshot"]>

package/src/storage/index.ts

@@ -16,6 +16,7 @@ export type {
   ProjectionExecution,
   ProjectionMaterializationIdentity,
 } from "../materialization/model"
+export type { ModelProviderIds } from "../models/events"
 export type {
   ActionRunEffectsRecord,
   ActionRunFailure,
@@ -65,6 +66,7 @@ export type {
   AgentRunExecution,
   AgentRunFailureCode,
   AgentRunFinishReason,
+  AgentRunKind,
   AgentRunRecord,
   AgentRunStatus,
   AgentRunStore,
@@ -75,9 +77,12 @@ export type {
   AgentThreadStore,
   AppendAgentMessageInput,
   ConfirmAgentRunExecutionOwnershipInput,
+  ConversationAgentRunRecord,
+  ConversationAgentRunSpec,
   CreateAgentContextCheckpointInput,
   CreateAgentRunInput,
   CreateAgentThreadInput,
+  CreateSubagentRunInput,
   DeleteAgentMessagesByRunInput,
   FinishAgentRunInput,
   FinishQueuedAgentRunInput,
@@ -89,6 +94,9 @@ export type {
   ListAgentThreadsResult,
   ReclaimAgentRunInput,
   StartAgentRunInput,
+  SubagentRunRecord,
+  SubagentRunResult,
+  SubagentRunSpec,
 } from "./agents"
 export {
   AGENT_RUN_DIAGNOSTIC_CODES,
@@ -111,11 +119,15 @@ export type {
   AiBillableMeter,
   AiBillingIdentity,
   AiCostComponent,
+  AiCostEstimate,
   AiCostStorage,
   AiCostStorageErrorCode,
   AiCostSummary,
   AiModelCallAccountingItem,
   AiModelCallCostRecord,
+  AiModelCallExecutionSummary,
+  AiModelCallGroup,
+  AiModelCallGroupSummary,
   AiMoney,
   AiPriceSource,
   AiPricingContext,
@@ -126,10 +138,42 @@ export type {
   InMemoryAiCostStorageSnapshot,
   ListAiModelCallAccountingInput,
   ListAiModelCallAccountingResult,
+  ListAiModelCallGroupsInput,
+  ListAiModelCallGroupsResult,
   QueryAiAccountingOverviewInput,
   SummarizeAiCostExecutionsInput,
 } from "./ai-cost"
 export { AiCostStorageError, InMemoryAiCostStorage } from "./ai-cost"
+export type {
+  AiLimitConsumption,
+  AiLimitMeter,
+  AiLimitPeriod,
+  AiLimitPeriodKind,
+  AiLimitPolicy,
+  AiLimitPolicyStatus,
+  AiLimitQuantity,
+  AiLimitReservationBucket,
+  AiLimitStorage,
+  AiLimitStorageErrorCode,
+  AiLimitSubject,
+  AiModelCallReservation,
+  AiModelCallReservationIdentity,
+  AiModelCallReservationState,
+  CreateAiLimitPolicyInput,
+  DeleteAiLimitPolicyInput,
+  GetAiLimitPolicyInput,
+  InMemoryAiLimitStorageOptions,
+  InMemoryAiLimitStorageSnapshot,
+  ListAiLimitPoliciesInput,
+  ListAiLimitPolicyStatusesInput,
+  MarkAiModelCallReservationUnknownInput,
+  ReconcileAiModelCallInput,
+  RecordAiModelCallLimitActualsInput,
+  ReserveAiModelCallInput,
+  ReserveAiModelCallResult,
+  UpdateAiLimitPolicyInput,
+} from "./ai-limits"
+export { AiLimitStorageError, aiLimitCalendarMonth, InMemoryAiLimitStorage } from "./ai-limits"
 export type {
   AiModelCallUsage,
   AiModelCallUsageInput,
@@ -158,6 +202,7 @@ export type {
   AccessTokenRecord,
   AccessTokenSubjectType,
   AuthAccessTokenStore,
+  AuthDeviceAuthorizationStore,
   AuthGroupMembershipStore,
   AuthInvitationStore,
   AuthMagicLinkStore,
@@ -170,6 +215,8 @@ export type {
   AuthUserIdentityStore,
   AuthUserStore,
   CompleteAuthSessionInput,
+  CompleteDeviceAuthorizationInput,
+  CompleteDeviceAuthorizationResult,
   CompleteMagicLinkSignInInput,
   CompleteOidcSignInInput,
   CompleteSignInResult,
@@ -178,8 +225,11 @@ export type {
   CreateAuthServiceAccountInput,
   CreateAuthSessionInput,
   CreateAuthUserInput,
+  CreateDeviceAuthorizationInput,
   CreateOidcAuthorizationAttemptInput,
   CreateOrUpdateAuthInvitationInput,
+  DeviceAuthorizationRecord,
+  DeviceAuthorizationStatus,
   GroupMembershipRecord,
   GroupMembershipSource,
   InvitationRecord,
@@ -210,7 +260,7 @@ export type {
   UserRecord,
   UserStatus,
 } from "./auth"
-export { AuthStorageError, InMemoryAuthStorage } from "./auth"
+export { AuthStorageError, InMemoryAuthStorage, MAX_PENDING_DEVICE_AUTHORIZATIONS } from "./auth"
 export type {
   AttachConnectorConnectionRunAuthorizationInput,
   ClaimConnectorConnectionRunCallbackInput,

package/src/storage/types.ts

@@ -1,6 +1,7 @@
 import type { ActionRunStorage } from "./action-runs"
 import type { AgentStorage } from "./agents"
 import type { AiCostStorage } from "./ai-cost"
+import type { AiLimitStorage } from "./ai-limits"
 import type { AiUsageStorage } from "./ai-usage"
 import type { AuthStorage } from "./auth"
 import type { ConnectorConnectionStorage } from "./connector-connections"
@@ -43,13 +44,17 @@ export type {
   AgentMessageRecord,
   AgentRunExecution,
   AgentRunFailureCode,
+  AgentRunKind,
   AgentRunRecord,
   AgentStorage,
   AgentThreadRecord,
   AppendAgentMessageInput,
   ConfirmAgentRunExecutionOwnershipInput,
+  ConversationAgentRunRecord,
+  ConversationAgentRunSpec,
   CreateAgentRunInput,
   CreateAgentThreadInput,
+  CreateSubagentRunInput,
   FinishAgentRunInput,
   FinishQueuedAgentRunInput,
   ListAgentMessagesInput,
@@ -60,8 +65,39 @@ export type {
   ListAgentThreadsResult,
   ReclaimAgentRunInput,
   StartAgentRunInput,
+  SubagentRunRecord,
+  SubagentRunResult,
+  SubagentRunSpec,
 } from "./agents"
 export { AGENT_RUN_FAILURE_CODES, AgentStorageError } from "./agents"
+export type {
+  AiLimitConsumption,
+  AiLimitMeter,
+  AiLimitPeriod,
+  AiLimitPeriodKind,
+  AiLimitPolicy,
+  AiLimitPolicyStatus,
+  AiLimitQuantity,
+  AiLimitReservationBucket,
+  AiLimitStorage,
+  AiLimitStorageErrorCode,
+  AiLimitSubject,
+  AiModelCallReservation,
+  AiModelCallReservationIdentity,
+  AiModelCallReservationState,
+  CreateAiLimitPolicyInput,
+  DeleteAiLimitPolicyInput,
+  GetAiLimitPolicyInput,
+  ListAiLimitPoliciesInput,
+  ListAiLimitPolicyStatusesInput,
+  MarkAiModelCallReservationUnknownInput,
+  ReconcileAiModelCallInput,
+  RecordAiModelCallLimitActualsInput,
+  ReserveAiModelCallInput,
+  ReserveAiModelCallResult,
+  UpdateAiLimitPolicyInput,
+} from "./ai-limits"
+export { AiLimitStorageError, aiLimitCalendarMonth } from "./ai-limits"
 export type {
   AiModelCallUsage,
   AiModelCallUsageInput,
@@ -295,6 +331,7 @@ export interface Storage {
   agents?: AgentStorage
   aiUsage?: AiUsageStorage
   aiCosts?: AiCostStorage
+  aiLimits?: AiLimitStorage
   actionRuns?: ActionRunStorage
   syncRuns?: SyncRunStorage
   pipelineRuns?: PipelineRunStorage

package/src/storage/workflow-runs/in-memory.ts

@@ -628,7 +628,7 @@ export class InMemoryWorkflowAgentNodeRunStorage implements WorkflowAgentNodeRun
       projectId: input.projectId,
       executionId: input.executionId,
       nodeRunId: input.nodeRunId,
-      agentId: input.agentId,
+      actorId: input.actorId,
       workflowExecutionId: workflowRun.executionId,
     })
     if (node.status !== "running") {
@@ -655,7 +655,7 @@ export class InMemoryWorkflowAgentNodeRunStorage implements WorkflowAgentNodeRun
       projectId: input.projectId,
       nodeRunId: input.nodeRunId,
       executionId: input.executionId,
-      agentId: input.agentId,
+      actorId: input.actorId,
       status: "queued",
       prompt: input.prompt,
       attempt: 0,
@@ -747,7 +747,7 @@ export class InMemoryWorkflowAgentNodeRunStorage implements WorkflowAgentNodeRun
     const order = input.order ?? "desc"
     const filtered = [...this.runs.values()]
       .filter((run) => run.projectId === input.projectId)
-      .filter((run) => (input.agentId ? run.agentId === input.agentId : true))
+      .filter((run) => (input.actorId ? run.actorId === input.actorId : true))
       .filter((run) => (statuses ? statuses.has(run.status) : true))
       .sort((a, b) =>
         order === "asc"

package/src/storage/workflow-runs/provider.ts

@@ -42,7 +42,7 @@ export async function assertWorkflowAgentNodeRunExecution(input: {
   readonly projectId: string
   readonly executionId: string
   readonly nodeRunId: string
-  readonly agentId: string
+  readonly actorId: string
   readonly workflowExecutionId: string
 }): Promise<void> {
   const execution = await findAgentRunExecution({
@@ -50,7 +50,7 @@ export async function assertWorkflowAgentNodeRunExecution(input: {
     projectId: input.projectId,
     executionId: input.executionId,
     runId: input.nodeRunId,
-    serviceAccountId: agentServiceAccountId(input.agentId),
+    authority: { type: "managed", serviceAccountId: agentServiceAccountId(input.actorId) },
   })
   if (
     !execution ||

package/src/storage/workflow-runs/types.ts

@@ -21,6 +21,8 @@ export const WORKFLOW_RUN_FAILURE_CODES = [
   "internal.unexpected",
   "runtime.cancelled",
   "workflow.node_failed",
+  "ai.usage_limit_exceeded",
+  "ai.usage_limit_unavailable",
 ] as const satisfies readonly [SixbErrorCode, ...SixbErrorCode[]]
 
 export type WorkflowRunFailureCode = (typeof WORKFLOW_RUN_FAILURE_CODES)[number]
@@ -42,7 +44,8 @@ export interface WorkflowAgentNodeRunRecord {
   readonly projectId: string
   readonly nodeRunId: string
   readonly executionId: string
-  readonly agentId: string
+  /** Stable managed actor for the workflow + Agent step, not the step definition id. */
+  readonly actorId: string
   readonly status: WorkflowAgentNodeRunStatus
   readonly prompt: string
   readonly modelId?: string
@@ -61,7 +64,7 @@ export interface CreateWorkflowAgentNodeRunInput {
   readonly projectId: string
   readonly nodeRunId: string
   readonly executionId: string
-  readonly agentId: string
+  readonly actorId: string
   readonly prompt: string
   readonly createdAt?: Date
 }
@@ -116,7 +119,7 @@ export interface CancelWorkflowAgentNodeRunInput {
 
 export interface ListWorkflowAgentNodeRunsInput {
   readonly projectId: string
-  readonly agentId?: string
+  readonly actorId?: string
   readonly statuses?: readonly WorkflowAgentNodeRunStatus[]
   readonly limit?: number
   readonly offset?: number

package/src/testing/agent-execution.ts

@@ -1,4 +1,5 @@
 import { agentServiceAccountId } from "../agents/authority"
+import { SYSTEM_PRINCIPAL } from "../auth"
 import type { AuthStorage } from "../storage/auth"
 import { AuthStorageError } from "../storage/auth"
 import type { ExecutionStorage } from "../storage/executions"
@@ -8,35 +9,43 @@ export async function createTestAgentExecution(
   storage: { readonly auth?: AuthStorage; readonly executions: ExecutionStorage },
   input: {
     readonly projectId: string
-    readonly agentId: string
+    readonly actorId?: string
     readonly runId: string
     readonly executionId?: string
     readonly sourceExecutionId?: string
+    /** Conversational runs inherit authority; workflow fixtures use managed authority. */
+    readonly authority?: "managed" | "inherited"
   }
 ): Promise<string> {
-  const auth = storage.auth
-  if (!auth) throw new Error("Agent execution fixtures require auth storage.")
-  const serviceAccountId = agentServiceAccountId(input.agentId)
-  const serviceAccount = await auth.serviceAccounts.getById({
-    projectId: input.projectId,
-    id: serviceAccountId,
-  })
-  if (!serviceAccount) {
-    const now = new Date()
-    try {
-      await auth.serviceAccounts.create({
-        id: serviceAccountId,
-        projectId: input.projectId,
-        name: `Test Agent ${input.agentId}`,
-        description: `Test service account for Agent '${input.agentId}'.`,
-        status: "active",
-        createdByPrincipal: { type: "system", id: "test" },
-        createdAt: now,
-        updatedAt: now,
-      })
-    } catch (error) {
-      if (!(error instanceof AuthStorageError) || error.code !== "duplicate_service_account") {
-        throw error
+  const authority = input.authority ?? (input.actorId === undefined ? "inherited" : "managed")
+  if (authority === "managed" && input.actorId === undefined) {
+    throw new Error("Managed Agent execution fixtures require an agent id.")
+  }
+  const serviceAccountId = input.actorId ? agentServiceAccountId(input.actorId) : undefined
+  if (authority === "managed") {
+    const auth = storage.auth
+    if (!auth) throw new Error("Managed Agent execution fixtures require auth storage.")
+    const serviceAccount = await auth.serviceAccounts.getById({
+      projectId: input.projectId,
+      id: serviceAccountId!,
+    })
+    if (!serviceAccount) {
+      const now = new Date()
+      try {
+        await auth.serviceAccounts.create({
+          id: serviceAccountId!,
+          projectId: input.projectId,
+          name: `Test Agent ${input.actorId}`,
+          description: `Test service account for Agent '${input.actorId}'.`,
+          status: "active",
+          createdByPrincipal: SYSTEM_PRINCIPAL,
+          createdAt: now,
+          updatedAt: now,
+        })
+      } catch (error) {
+        if (!(error instanceof AuthStorageError) || error.code !== "duplicate_service_account") {
+          throw error
+        }
       }
     }
   }
@@ -72,10 +81,13 @@ export async function createTestAgentExecution(
       ? {}
       : { requestedBy: structuredClone(parent.requestedBy) }),
     correlationId: parent.correlationId,
-    authorizationRef: {
-      type: "principal",
-      principal: { type: "serviceAccount", id: serviceAccountId },
-    },
+    authorizationRef:
+      authority === "inherited"
+        ? structuredClone(parent.authorizationRef)
+        : {
+            type: "principal",
+            principal: { type: "serviceAccount", id: serviceAccountId! },
+          },
   })
   return executionId
 }