From bec0e9201d86c756db8cb7e2dc4e19de3abf7638 Mon Sep 17 00:00:00 2001 From: Saurabh Jain Date: Sun, 1 Mar 2026 16:08:23 +0530 Subject: [PATCH] feat: add trace_id and ToolContext to workflow types --- .../java/com/getaxonflow/sdk/AxonFlow.java | 3 + .../sdk/types/workflow/WorkflowTypes.java | 131 ++++++++++++++++-- 2 files changed, 126 insertions(+), 8 deletions(-) diff --git a/src/main/java/com/getaxonflow/sdk/AxonFlow.java b/src/main/java/com/getaxonflow/sdk/AxonFlow.java index 3c371bb..f937af8 100644 --- a/src/main/java/com/getaxonflow/sdk/AxonFlow.java +++ b/src/main/java/com/getaxonflow/sdk/AxonFlow.java @@ -4303,6 +4303,9 @@ public com.getaxonflow.sdk.types.workflow.WorkflowTypes.ListWorkflowsResponse li if (options.getOffset() > 0) { appendQueryParam(query, "offset", String.valueOf(options.getOffset())); } + if (options.getTraceId() != null) { + appendQueryParam(query, "trace_id", options.getTraceId()); + } } if (query.length() > 0) { diff --git a/src/main/java/com/getaxonflow/sdk/types/workflow/WorkflowTypes.java b/src/main/java/com/getaxonflow/sdk/types/workflow/WorkflowTypes.java index a68ba75..3ebc2e6 100644 --- a/src/main/java/com/getaxonflow/sdk/types/workflow/WorkflowTypes.java +++ b/src/main/java/com/getaxonflow/sdk/types/workflow/WorkflowTypes.java @@ -23,6 +23,7 @@ import java.time.Instant; import java.util.Collections; +import java.util.HashMap; import java.util.List; import java.util.Map; import java.util.Objects; @@ -230,16 +231,21 @@ public static final class CreateWorkflowRequest { @JsonProperty("metadata") private final Map metadata; + @JsonProperty("trace_id") + private final String traceId; + @JsonCreator public CreateWorkflowRequest( @JsonProperty("workflow_name") String workflowName, @JsonProperty("source") WorkflowSource source, @JsonProperty("total_steps") Integer totalSteps, - @JsonProperty("metadata") Map metadata) { + @JsonProperty("metadata") Map metadata, + @JsonProperty("trace_id") String traceId) { this.workflowName = Objects.requireNonNull(workflowName, "workflowName is required"); this.source = source != null ? source : WorkflowSource.EXTERNAL; this.totalSteps = totalSteps; this.metadata = metadata != null ? Collections.unmodifiableMap(metadata) : Collections.emptyMap(); + this.traceId = traceId; } public String getWorkflowName() { @@ -258,6 +264,10 @@ public Map getMetadata() { return metadata; } + public String getTraceId() { + return traceId; + } + public static Builder builder() { return new Builder(); } @@ -267,6 +277,7 @@ public static final class Builder { private WorkflowSource source = WorkflowSource.EXTERNAL; private Integer totalSteps; private Map metadata; + private String traceId; public Builder workflowName(String workflowName) { this.workflowName = workflowName; @@ -288,8 +299,13 @@ public Builder metadata(Map metadata) { return this; } + public Builder traceId(String traceId) { + this.traceId = traceId; + return this; + } + public CreateWorkflowRequest build() { - return new CreateWorkflowRequest(workflowName, source, totalSteps, metadata); + return new CreateWorkflowRequest(workflowName, source, totalSteps, metadata, traceId); } } } @@ -315,18 +331,23 @@ public static final class CreateWorkflowResponse { @JsonProperty("created_at") private final Instant createdAt; + @JsonProperty("trace_id") + private final String traceId; + @JsonCreator public CreateWorkflowResponse( @JsonProperty("workflow_id") String workflowId, @JsonProperty("workflow_name") String workflowName, @JsonProperty("source") WorkflowSource source, @JsonProperty("status") WorkflowStatus status, - @JsonProperty("created_at") Instant createdAt) { + @JsonProperty("created_at") Instant createdAt, + @JsonProperty("trace_id") String traceId) { this.workflowId = workflowId; this.workflowName = workflowName; this.source = source; this.status = status; this.createdAt = createdAt; + this.traceId = traceId; } public String getWorkflowId() { @@ -348,6 +369,64 @@ public WorkflowStatus getStatus() { public Instant getCreatedAt() { return createdAt; } + + public String getTraceId() { + return traceId; + } + } + + /** + * Tool-level context for per-tool governance within tool_call steps. + */ + @JsonIgnoreProperties(ignoreUnknown = true) + public static final class ToolContext { + + @JsonProperty("tool_name") + private final String toolName; + + @JsonProperty("tool_type") + private final String toolType; + + @JsonProperty("tool_input") + private final Map toolInput; + + private ToolContext(Builder builder) { + this.toolName = builder.toolName; + this.toolType = builder.toolType; + this.toolInput = builder.toolInput != null ? Collections.unmodifiableMap(new HashMap<>(builder.toolInput)) : null; + } + + @JsonCreator + public ToolContext( + @JsonProperty("tool_name") String toolName, + @JsonProperty("tool_type") String toolType, + @JsonProperty("tool_input") Map toolInput) { + this.toolName = toolName; + this.toolType = toolType; + this.toolInput = toolInput != null ? Collections.unmodifiableMap(new HashMap<>(toolInput)) : null; + } + + public String getToolName() { return toolName; } + public String getToolType() { return toolType; } + public Map getToolInput() { return toolInput; } + + public static Builder builder(String toolName) { + return new Builder(toolName); + } + + public static final class Builder { + private final String toolName; + private String toolType; + private Map toolInput; + + public Builder(String toolName) { + this.toolName = Objects.requireNonNull(toolName, "toolName must not be null"); + } + + public Builder toolType(String toolType) { this.toolType = toolType; return this; } + public Builder toolInput(Map toolInput) { this.toolInput = toolInput; return this; } + public ToolContext build() { return new ToolContext(this); } + } } /** @@ -371,18 +450,23 @@ public static final class StepGateRequest { @JsonProperty("provider") private final String provider; + @JsonProperty("tool_context") + private final ToolContext toolContext; + @JsonCreator public StepGateRequest( @JsonProperty("step_name") String stepName, @JsonProperty("step_type") StepType stepType, @JsonProperty("step_input") Map stepInput, @JsonProperty("model") String model, - @JsonProperty("provider") String provider) { + @JsonProperty("provider") String provider, + @JsonProperty("tool_context") ToolContext toolContext) { this.stepName = stepName; this.stepType = Objects.requireNonNull(stepType, "stepType is required"); this.stepInput = stepInput != null ? Collections.unmodifiableMap(stepInput) : Collections.emptyMap(); this.model = model; this.provider = provider; + this.toolContext = toolContext; } public String getStepName() { @@ -405,6 +489,10 @@ public String getProvider() { return provider; } + public ToolContext getToolContext() { + return toolContext; + } + public static Builder builder() { return new Builder(); } @@ -415,6 +503,7 @@ public static final class Builder { private Map stepInput; private String model; private String provider; + private ToolContext toolContext; public Builder stepName(String stepName) { this.stepName = stepName; @@ -441,8 +530,13 @@ public Builder provider(String provider) { return this; } + public Builder toolContext(ToolContext toolContext) { + this.toolContext = toolContext; + return this; + } + public StepGateRequest build() { - return new StepGateRequest(stepName, stepType, stepInput, model, provider); + return new StepGateRequest(stepName, stepType, stepInput, model, provider, toolContext); } } } @@ -679,6 +773,9 @@ public static final class WorkflowStatusResponse { @JsonProperty("steps") private final List steps; + @JsonProperty("trace_id") + private final String traceId; + @JsonCreator public WorkflowStatusResponse( @JsonProperty("workflow_id") String workflowId, @@ -689,7 +786,8 @@ public WorkflowStatusResponse( @JsonProperty("total_steps") Integer totalSteps, @JsonProperty("started_at") Instant startedAt, @JsonProperty("completed_at") Instant completedAt, - @JsonProperty("steps") List steps) { + @JsonProperty("steps") List steps, + @JsonProperty("trace_id") String traceId) { this.workflowId = workflowId; this.workflowName = workflowName; this.source = source; @@ -699,6 +797,7 @@ public WorkflowStatusResponse( this.startedAt = startedAt; this.completedAt = completedAt; this.steps = steps != null ? Collections.unmodifiableList(steps) : Collections.emptyList(); + this.traceId = traceId; } public String getWorkflowId() { @@ -737,6 +836,10 @@ public List getSteps() { return steps; } + public String getTraceId() { + return traceId; + } + public boolean isTerminal() { return status == WorkflowStatus.COMPLETED || status == WorkflowStatus.ABORTED || @@ -753,12 +856,14 @@ public static final class ListWorkflowsOptions { private final WorkflowSource source; private final int limit; private final int offset; + private final String traceId; - public ListWorkflowsOptions(WorkflowStatus status, WorkflowSource source, int limit, int offset) { + public ListWorkflowsOptions(WorkflowStatus status, WorkflowSource source, int limit, int offset, String traceId) { this.status = status; this.source = source; this.limit = limit > 0 ? limit : 50; this.offset = Math.max(offset, 0); + this.traceId = traceId; } public WorkflowStatus getStatus() { @@ -777,6 +882,10 @@ public int getOffset() { return offset; } + public String getTraceId() { + return traceId; + } + public static Builder builder() { return new Builder(); } @@ -786,6 +895,7 @@ public static final class Builder { private WorkflowSource source; private int limit = 50; private int offset = 0; + private String traceId; public Builder status(WorkflowStatus status) { this.status = status; @@ -807,8 +917,13 @@ public Builder offset(int offset) { return this; } + public Builder traceId(String traceId) { + this.traceId = traceId; + return this; + } + public ListWorkflowsOptions build() { - return new ListWorkflowsOptions(status, source, limit, offset); + return new ListWorkflowsOptions(status, source, limit, offset, traceId); } } }