Return Guardrail token usage (#52)

* Return Guardrail token usage
* Return token count even during execution error

Author: steven10a

docs/agents_sdk_integration.md

@@ -111,3 +111,28 @@ const agent = await GuardrailAgent.create(
 - Explore available guardrails for your use case  
 - Learn about pipeline configuration in our [quickstart](./quickstart.md)
 - For more details on the OpenAI Agents SDK, refer to the [Agent SDK documentation](https://openai.github.io/openai-agents-js/).
+
+## Token Usage Tracking
+
+!!! warning "JavaScript Agents SDK Limitation"
+    The JavaScript Agents SDK (`@openai/agents`) does not currently return guardrail results in the `RunResult` object. This means `totalGuardrailTokenUsage()` cannot retrieve token counts from Agents SDK runs.
+    
+    **For token usage tracking, use `GuardrailsOpenAI` instead of `GuardrailAgent`.** The Python Agents SDK does support this feature.
+
+When a guardrail **triggers** (throws `InputGuardrailTripwireTriggered` or `OutputGuardrailTripwireTriggered`), token usage IS available in the error's result object:
+
+```typescript
+try {
+  const result = await Runner.run(agent, userInput);
+} catch (error) {
+  if (error.constructor.name === 'InputGuardrailTripwireTriggered') {
+    // Token usage available when guardrail triggers
+    const usage = error.result?.output?.outputInfo?.token_usage;
+    if (usage) {
+      console.log(`Guardrail tokens: ${usage.total_tokens}`);
+    }
+  }
+}
+```
+
+For full token usage tracking across all guardrail runs (passing and failing), use the `GuardrailsOpenAI` client instead - see the [quickstart](./quickstart.md#token-usage-tracking) for details.

docs/quickstart.md

@@ -205,6 +205,75 @@ const client = await GuardrailsOpenAI.create(
 );
 ```
 
+## Token Usage Tracking
+
+LLM-based guardrails (Jailbreak, custom prompt checks, etc.) consume tokens. Keep track of those costs with the `totalGuardrailTokenUsage` helper:
+
+```typescript
+import { GuardrailsOpenAI, totalGuardrailTokenUsage } from '@openai/guardrails';
+
+const client = await GuardrailsOpenAI.create(CONFIG);
+const response = await client.guardrails.responses.create({
+  model: 'gpt-4.1-mini',
+  input: 'Hello!',
+});
+
+const tokens = totalGuardrailTokenUsage(response);
+console.log(`Guardrail tokens used: ${tokens.total_tokens}`);
+// => Guardrail tokens used: 425
+```
+
+The helper returns:
+
+```typescript
+{
+  prompt_tokens: 300,       // Sum of prompt tokens across all LLM guardrails
+  completion_tokens: 125,   // Sum of completion tokens
+  total_tokens: 425,        // Total guardrail tokens
+}
+```
+
+### Works With GuardrailsOpenAI Clients
+
+`totalGuardrailTokenUsage` works across all client types and endpoints:
+
+- **OpenAI** - sync and async clients
+- **Azure OpenAI** - sync and async clients  
+- **Third-party providers** - any OpenAI-compatible API wrapper
+- **Endpoints** - both `responses` and `chat.completions`
+- **Streaming** - capture from the final chunk
+
+```typescript
+// OpenAI client responses
+const response = await client.guardrails.responses.create(...);
+const tokens = totalGuardrailTokenUsage(response);
+
+// Streaming – use the final chunk
+let lastChunk: unknown;
+for await (const chunk of stream) {
+  lastChunk = chunk;
+}
+const streamingTokens = lastChunk ? totalGuardrailTokenUsage(lastChunk) : null;
+```
+
+**Note:** The JavaScript Agents SDK (`@openai/agents`) does not currently populate guardrail results in the `RunResult` object, so `totalGuardrailTokenUsage()` will return empty results for Agents SDK runs.
+
+### Per-Guardrail Usage
+
+Each guardrail result includes its own `token_usage` entry:
+
+```typescript
+const response = await client.guardrails.responses.create(...);
+for (const gr of response.guardrail_results.allResults) {
+  const usage = gr.info.token_usage;
+  if (usage) {
+    console.log(`${gr.info.guardrail_name}: ${usage.total_tokens} tokens`);
+  }
+}
+```
+
+Non-LLM guardrails (PII, Moderation, URL Filter, etc.) do not consume tokens, so `token_usage` will be omitted.
+
 ## Next Steps
 
 - Explore TypeScript [examples](https://github.com/openai/openai-guardrails-js/tree/main/examples) for advanced patterns

examples/basic/hello_world.ts

@@ -8,19 +8,23 @@
  */
 
 import * as readline from 'readline';
-import { GuardrailsOpenAI, GuardrailTripwireTriggered } from '../../src';
+import { GuardrailsOpenAI, GuardrailTripwireTriggered, totalGuardrailTokenUsage } from '../../src';
 
-// Pipeline configuration with preflight PII masking and input guardrails
+// Pipeline configuration with preflight and input guardrails
 const PIPELINE_CONFIG = {
   version: 1,
   pre_flight: {
     version: 1,
     guardrails: [
       {
-        name: 'Contains PII',
+        name: 'Moderation',
+        config: { categories: ['hate', 'violence'] },
+      },
+      {
+        name: 'Jailbreak',
         config: {
-          entities: ['US_SSN', 'PHONE_NUMBER', 'EMAIL_ADDRESS'],
-          block: true, // Use masking mode (default) - masks PII without blocking
+          model: 'gpt-4.1-mini',
+          confidence_threshold: 0.7,
         },
       },
     ],
@@ -76,6 +80,10 @@ async function processInput(
   // Show guardrail results if any were run
   if (response.guardrail_results.allResults.length > 0) {
     console.log(`[dim]Guardrails checked: ${response.guardrail_results.allResults.length}[/dim]`);
+    const usage = totalGuardrailTokenUsage(response);
+    if (usage.total_tokens !== null) {
+      console.log(`[dim]Token usage: ${JSON.stringify(usage)}[/dim]`);
+    }
   }
 
   return response.id;

examples/basic/local_model.ts

@@ -2,7 +2,7 @@
  * Example: Guardrail bundle using Ollama's Gemma3 model with GuardrailsClient.
  */
 
-import { GuardrailsOpenAI, GuardrailTripwireTriggered } from '../../src';
+import { GuardrailsOpenAI, GuardrailTripwireTriggered, totalGuardrailTokenUsage } from '../../src';
 import * as readline from 'readline';
 import { OpenAI } from 'openai';
 
@@ -46,6 +46,10 @@ async function processInput(
     // Access response content using standard OpenAI API
     const responseContent = response.choices[0].message.content ?? '';
     console.log(`\nAssistant output: ${responseContent}\n`);
+    const usage = totalGuardrailTokenUsage(response);
+    if (usage.total_tokens !== null) {
+      console.log(`Token usage: ${usage.total_tokens}`);
+    }
 
     // Guardrails passed - now safe to add to conversation history
     conversation.push({ role: 'user', content: userInput });

examples/basic/multiturn_with_prompt_injection_detection.ts

@@ -26,7 +26,12 @@
  */
 
 import * as readline from 'readline';
-import { GuardrailsOpenAI, GuardrailTripwireTriggered, GuardrailsResponse } from '../../src';
+import {
+  GuardrailsOpenAI,
+  GuardrailTripwireTriggered,
+  GuardrailsResponse,
+  totalGuardrailTokenUsage,
+} from '../../src';
 
 // Tool implementations (mocked)
 function get_horoscope(sign: string): { horoscope: string } {
@@ -299,6 +304,15 @@ async function main(malicious: boolean = false): Promise<void> {
 
         printGuardrailResults('initial', response);
 
+        const initialUsage = totalGuardrailTokenUsage(response);
+        if (initialUsage.total_tokens !== null) {
+          console.log(
+            `[dim]Guardrail tokens (initial): ${initialUsage.total_tokens} · prompt=${
+              initialUsage.prompt_tokens ?? 0
+            }, completion=${initialUsage.completion_tokens ?? 0}[/dim]`
+          );
+        }
+
         assistantOutputs = response.output ?? [];
 
         // Guardrails passed - now safe to add user message to conversation history
@@ -394,6 +408,14 @@ async function main(malicious: boolean = false): Promise<void> {
           });
 
           printGuardrailResults('final', response);
+          const finalUsage = totalGuardrailTokenUsage(response);
+          if (finalUsage.total_tokens !== null) {
+            console.log(
+              `[dim]Guardrail tokens (final): ${finalUsage.total_tokens} · prompt=${
+                finalUsage.prompt_tokens ?? 0
+              }, completion=${finalUsage.completion_tokens ?? 0}[/dim]`
+            );
+          }
           console.log(`\n🤖 Assistant: ${response.output_text}`);
 
           // Guardrails passed - now safe to add tool results and assistant responses to history

examples/basic/streaming.ts

@@ -3,7 +3,7 @@
  * Streams output using console logging.
  */
 
-import { GuardrailsOpenAI, GuardrailTripwireTriggered } from '../../src';
+import { GuardrailsOpenAI, GuardrailTripwireTriggered, totalGuardrailTokenUsage } from '../../src';
 import * as readline from 'readline';
 
 // Define your pipeline configuration
@@ -14,10 +14,14 @@ const PIPELINE_CONFIG = {
     version: 1,
     guardrails: [
       {
-        name: 'Contains PII',
+        name: 'Moderation',
+        config: { categories: ['hate', 'violence'] },
+      },
+      {
+        name: 'Jailbreak',
         config: {
-          entities: ['US_SSN', 'PHONE_NUMBER', 'EMAIL_ADDRESS'],
-          block: false, // Use masking mode (default) - masks PII without blocking
+          model: 'gpt-4.1-mini',
+          confidence_threshold: 0.7,
         },
       },
     ],
@@ -49,6 +53,7 @@ const PIPELINE_CONFIG = {
         config: {
           entities: ['US_SSN', 'PHONE_NUMBER', 'EMAIL_ADDRESS'],
           block: true, // Use blocking mode on output
+          detect_encoded_pii: false,
         },
       },
     ],
@@ -78,8 +83,10 @@ async function processInput(
     console.log(outputText);
 
     let responseIdToReturn: string | null = null;
+    let lastChunk: unknown = null;
 
     for await (const chunk of stream) {
+      lastChunk = chunk;
       // Access streaming response exactly like native OpenAI API
       if ('delta' in chunk && chunk.delta && typeof chunk.delta === 'string') {
         outputText += chunk.delta;
@@ -99,6 +106,14 @@ async function processInput(
     }
 
     console.log(); // New line after streaming
+
+    if (lastChunk) {
+      const usage = totalGuardrailTokenUsage(lastChunk);
+      if (usage.total_tokens !== null) {
+        console.log(`[dim]📊 Guardrail tokens: ${usage.total_tokens}[/dim]`);
+      }
+    }
+
     return responseIdToReturn;
   } catch (error) {
     if (error instanceof GuardrailTripwireTriggered) {

src/__tests__/unit/agents.test.ts

@@ -547,4 +547,68 @@ describe('GuardrailAgent', () => {
       );
     });
   });
+
+  it('propagates guardrail metadata to outputInfo on success', async () => {
+    process.env.OPENAI_API_KEY = 'test';
+    const config = {
+      version: 1,
+      input: {
+        version: 1,
+        guardrails: [{ name: 'Jailbreak', config: {} }],
+      },
+    };
+
+    const { instantiateGuardrails } = await import('../../runtime');
+    vi.mocked(instantiateGuardrails).mockImplementationOnce(() =>
+      Promise.resolve([
+        {
+          definition: {
+            name: 'Jailbreak',
+            description: 'Test guardrail',
+            mediaType: 'text/plain',
+            configSchema: z.object({}),
+            checkFn: vi.fn(),
+            metadata: {},
+            ctxRequirements: z.object({}),
+            schema: () => ({}),
+            instantiate: vi.fn(),
+          },
+          config: {},
+          run: vi.fn().mockResolvedValue({
+            tripwireTriggered: false,
+            info: {
+              guardrail_name: 'Jailbreak',
+              flagged: false,
+              token_usage: {
+                prompt_tokens: 42,
+                completion_tokens: 10,
+                total_tokens: 52,
+              },
+            },
+          }),
+        } as unknown as Parameters<typeof instantiateGuardrails>[0] extends Promise<infer T>
+          ? T extends readonly (infer U)[]
+            ? U
+            : never
+          : never,
+      ])
+    );
+
+    const agent = (await GuardrailAgent.create(
+      config,
+      'Metadata Agent',
+      'Test instructions'
+    )) as MockAgent;
+
+    const guardrailFunction = agent.inputGuardrails[0];
+    const result = await guardrailFunction.execute('payload');
+
+    expect(result.tripwireTriggered).toBe(false);
+    expect(result.outputInfo.guardrail_name).toBe('Jailbreak');
+    expect(result.outputInfo.token_usage).toEqual({
+      prompt_tokens: 42,
+      completion_tokens: 10,
+      total_tokens: 52,
+    });
+  });
 });