diff --git a/agent.py b/agent.py index c5da9b8..7527d20 100644 --- a/agent.py +++ b/agent.py @@ -342,7 +342,7 @@ class AgentComplianceAuditor(A2AAgent[AgentComplianceAuditorConfig, NoAuth]): await _event(ctx, "audit_completed", audit_id=audit_id, skill="audit_runtime", findings=len(findings)) return FindingResult(audit_id=audit_id, target=request.target, status="ok", output_dir=output_dir, generated_files=files, findings=findings, evidence=evidence, warnings=warnings) - @a2a.tool(description="Deduplicate findings, map evidence to controls, and score severity, confidence, and aggregate risk", timeout_seconds=120, idempotent=True, cost_class="read-only", input_schema=_tool_schema("risk")) + @a2a.tool(description="Deduplicate findings, map evidence to controls, and score severity, confidence, and aggregate risk", timeout_seconds=120, idempotent=True, cost_class="read-only") async def assess_risk(self, ctx: RunContext[NoAuth], request: AssessRiskRequest) -> RiskSummaryResult: audit_id = _audit_id(request.target, request.options.audit_id) await _event(ctx, "audit_started", audit_id=audit_id, skill="assess_risk", target=request.target.agent_name) @@ -353,7 +353,7 @@ class AgentComplianceAuditor(A2AAgent[AgentComplianceAuditorConfig, NoAuth]): await _event(ctx, "audit_completed", audit_id=audit_id, skill="assess_risk", findings=len(all_findings), score=summary["risk_score"]) return RiskSummaryResult(audit_id=audit_id, target=request.target, status="ok", output_dir=output_dir, generated_files=files, risk_summary=summary, findings=all_findings) - @a2a.tool(description="Generate actionable remediation and a machine-readable SARIF report for compliance findings", timeout_seconds=120, idempotent=True, cost_class="read-only") + @a2a.tool(description="Generate actionable remediation and a machine-readable SARIF report for compliance findings", timeout_seconds=120, idempotent=True, cost_class="read-only", input_schema=_tool_schema("remediation")) async def generate_remediation(self, ctx: RunContext[NoAuth], request: GenerateRemediationRequest) -> RemediationResult: audit_id = _audit_id(request.target, request.options.audit_id) await _event(ctx, "audit_started", audit_id=audit_id, skill="generate_remediation", target=request.target.agent_name)